import gradio as gr
import requests
import time
import os
from transformers import pipeline, set_seed

# 设置随机种子
set_seed(42)

# 从环境变量获取 token
HF_TOKEN = os.environ.get('HF_TOKEN')
token_status = "✅ 已设置" if HF_TOKEN else "❌ 未设置"

# 初始化模型管道
try:
    protein_generator = pipeline(
        "text-generation", 
        model="mzcwd/ProtTeX",
        device=-1,  # 使用CPU
        torch_dtype="auto"
    )
    MODEL_LOADED = True
    model_status = "✅ 本地模型加载成功"
    print(model_status)
except Exception as e:
    print(f"❌ 本地模型加载失败: {e}")
    MODEL_LOADED = False
    protein_generator = None
    model_status = "❌ 本地模型加载失败，使用API模式"

# 确定运行模式
if MODEL_LOADED:
    run_mode = "🧠 本地模型"
else:
    run_mode = "🌐 API调用"

def generate_with_local_model(instruction, max_length=100):
    """使用本地加载的模型生成蛋白质序列"""
    try:
        result = protein_generator(
            instruction,
            max_length=max_length,
            num_return_sequences=1,
            temperature=0.8,
            do_sample=True,
            pad_token_id=50256
        )
        
        if result and len(result) > 0:
            return f"✅ 生成成功:\n\n{result[0]['generated_text']}"
        else:
            return "❌ 生成失败，未获得有效结果"
            
    except Exception as e:
        return f"❌ 生成过程中出现错误: {str(e)}"

def generate_with_api(instruction, max_length=100):
    """使用HuggingFace API生成蛋白质序列"""
    if not HF_TOKEN:
        return "❌ 未设置 HuggingFace Token，无法使用 API 功能"
    
    try:
        API_URL = "https://api-inference.huggingface.co/models/mzcwd/ProtTeX"
        
        headers = {
            "Authorization": f"Bearer {HF_TOKEN}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "inputs": instruction,
            "parameters": {
                "max_new_tokens": max_length,
                "temperature": 0.8,
                "do_sample": True,
                "return_full_text": False
            }
        }
        
        response = requests.post(API_URL, headers=headers, json=payload, timeout=60)
        
        if response.status_code == 200:
            result = response.json()
            if isinstance(result, list) and len(result) > 0:
                generated_text = result[0].get('generated_text', '生成成功但无内容返回')
                return f"✅ API 生成成功:\n\n{generated_text}"
            return f"✅ API 响应: {str(result)}"
        else:
            error_msg = f"❌ API请求失败 (状态码: {response.status_code})"
            if response.text:
                error_msg += f"\n详细信息: {response.text}"
            return error_msg
            
    except requests.exceptions.Timeout:
        return "⏰ 请求超时，请稍后重试"
    except Exception as e:
        return f"❌ API调用错误: {str(e)}"

def generate_protein(instruction, max_length=100):
    """主生成函数"""
    if not instruction or instruction.strip() == "":
        return "❌ 请输入有效的蛋白质生成指令"
    
    # 显示处理状态
    time.sleep(0.5)
    
    try:
        # 优先使用本地模型，失败时使用API
        if MODEL_LOADED:
            result = generate_with_local_model(instruction, max_length)
        else:
            result = generate_with_api(instruction, max_length)
        
        return result
        
    except Exception as e:
        return f"❌ 生成过程中出现未预期错误: {str(e)}"

# 创建界面
with gr.Blocks(title="ProtTeX 蛋白质生成器") as demo:
    
    gr.Markdown(f"""
    # 🧬 ProtTeX 蛋白质生成器
    
    **使用自然语言指令生成蛋白质序列**
    
    *当前模式: {run_mode} | Token状态: {token_status}*
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            instruction = gr.Textbox(
                label="🧪 蛋白质生成指令",
                placeholder="例如：Generate a protein with alpha-helical structure for membrane binding",
                lines=3,
                value="Generate a hydrophobic transmembrane protein sequence"
            )
            
            max_length = gr.Slider(
                minimum=50,
                maximum=300,
                value=150,
                step=10,
                label="📏 序列最大长度"
            )
            
            generate_btn = gr.Button(
                "🚀 生成蛋白质序列",
                variant="primary"
            )
            
        with gr.Column(scale=1):
            output = gr.Textbox(
                label="🧬 生成的蛋白质序列",
                lines=8
            )
    
    # 示例部分
    gr.Markdown("### 💡 示例指令（点击试用）")
    
    examples = gr.Examples(
        examples=[
            ["Generate a hydrophobic transmembrane protein sequence"],
            ["Create a water-soluble protein with beta-sheet structure"],
            ["Design a protein with enzymatic activity for hydrolysis"],
            ["Generate a stable protein for high temperature environments"]
        ],
        inputs=[instruction],
        outputs=[output],
        fn=generate_protein,
        cache_examples=False
    )
    
    # 连接按钮事件
    generate_btn.click(
        fn=generate_protein,
        inputs=[instruction, max_length],
        outputs=[output]
    )
    
    # 状态信息
    gr.Markdown(f"""
    ---
    **系统状态**:
    - 运行模式: {run_mode}
    - Token状态: {token_status}
    - 模型状态: {model_status}
    - 硬件: CPU Basic
    
    *基于 [mzcwd/ProtTeX](https://huggingface.co/mzcwd/ProtTeX) 模型*
    """)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        share=False
    )