import os

os.system('pip install transformers -U')
os.system('pip install modelscope -U')
os.system('pip install accelerate')
from threading import Thread
from typing import Iterator

import gradio as gr
import torch
from modelscope import AutoModelForCausalLM, AutoTokenizer
from transformers import TextIteratorStreamer

MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))

# 系统提示
SYSTEM_PROMPT = """
作为助手，你的角色是通过系统的长期思考过程，深入探讨问题，然后提供最终精确的解决方案。
这需要通过分析、总结、探索、重新评估、反思、回溯和迭代等全面的思维循环来发展深思熟虑的思考过程。
请将你的回答结构化为两个主要部分：思考和解决方案。
在“思考”部分，使用指定的格式详细描述你的推理过程：
<think> {思考过程，每个步骤之间用“\n\n”分隔} </think>

每个步骤应包括详细的考虑事项，例如分析问题、总结相关发现、头脑风暴新想法、验证当前步骤的准确性、修正任何错误以及回顾之前的步骤。
在“解决方案”部分，根据“思考”部分的各种尝试、探索和反思，系统地呈现你认为正确的最终解决方案。
解决方案应保持逻辑准确、简洁表达，并详细说明达到结论所需的必要步骤，格式如下：
<answer> {最终格式化的、精确和清晰的解决方案} </answer>
"""

# 无论 GPU 是否可用，都加载模型和分词器
model_id = "daniel0527/qwen2.5-3b-dsitill"
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    device_map="auto" if device == "cuda" else None
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.use_default_system_prompt = False


def generate(
        message: str,
        chat_history: list[tuple[str, str]],
        system_prompt: str = SYSTEM_PROMPT,
        max_new_tokens: int = 1024,
        temperature: float = 0.6,
        top_p: float = 0.9,
        top_k: int = 50,
        repetition_penalty: float = 1.2,
) -> Iterator[str]:
    conversation = []
    if system_prompt:
        conversation.append({"role": "system", "content": system_prompt})
    for user, assistant in chat_history:
        conversation.extend([
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant}
        ])
    conversation.append({"role": "user", "content": message})

    # 生成对话模板字符串，并进行编码时添加 padding 以获得 attention_mask
    input_str = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([input_str], return_tensors="pt", padding=True)
    inputs = inputs.to(model.device)

    # 调整 timeout 为 30 秒，必要时也可增加
    streamer = TextIteratorStreamer(tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        input_ids=inputs.input_ids,
        attention_mask=inputs.attention_mask,  # 传入 attention_mask
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        repetition_penalty=repetition_penalty,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    # 使用 try/except 捕获可能的 Empty 异常，确保生成过程结束后退出
    try:
        for text in streamer:
            outputs.append(text)
            yield "".join(outputs)
    except Exception:
        yield "".join(outputs)


# 创建 Gradio ChatInterface，不在此处添加示例问题
chat_interface = gr.ChatInterface(
    fn=generate,
    additional_inputs=[
        gr.Textbox(label="系统提示", value=SYSTEM_PROMPT, lines=6),
        gr.Slider(
            label="最大生成标记数",
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=DEFAULT_MAX_NEW_TOKENS,
        ),
        gr.Slider(
            label="温度",
            minimum=0.1,
            maximum=4.0,
            step=0.1,
            value=0.6,
        ),
        gr.Slider(
            label="Top-p（核采样）",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.9,
        ),
        gr.Slider(
            label="Top-k",
            minimum=1,
            maximum=1000,
            step=1,
            value=50,
        ),
        gr.Slider(
            label="重复惩罚",
            minimum=1.0,
            maximum=2.0,
            step=0.05,
            value=1.2,
        ),
    ],
    stop_btn="停止",
)

# 定义示例问题列表（将在聊天框下方展示）
example_questions = [
    "在患者做闭眼动作时，患侧眼球向外上方转动的现象被称为什么？",
    "对于一位病情稳定且血压在正常范围内波动，且未发生类似发作的患者，为了帮助诊断，应该进行哪种刺激试验？?",
    "肚子疼去那个科室？",
]

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(
        """<p align="center"><img src="https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png" style="height: 80px"/></p>""")
    gr.Markdown("""<center><font size=8>jishi-3b-r1-Chat Bot👾</center>""")
    gr.Markdown("""<center><font size=4>jishi-3b-r1是医疗行业30亿规模的大模型。</center>""")

    chat_interface.render()

    # 在聊天组件下方添加示例问题
    gr.Markdown("### 示例问题")
    for question in example_questions:
        gr.Markdown(f"- {question}")

if __name__ == "__main__":
    # 如果需要创建公网链接，请设置 share=True
    demo.queue(max_size=20).launch()