import os
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# ========= 基本設定 =========
# 你的模型 repo id（目前是 private 也沒關係）
MODEL_ID = "aciang/mistral7b-tk-sft-20251019-merged"

# 若模型是 private，建議在 Space 的「Settings → Repository secrets」加上 HF_TOKEN
HF_TOKEN = os.getenv("HF_TOKEN", None)

# 建議預設的「傳統知識」系統提示，可以在介面中修改
DEFAULT_SYSTEM_PROMPT = """你是一位熟悉台灣與國際原住民族傳統知識的學者，
擅長用淺顯但尊重文化脈絡的繁體中文說明各族的夢境、儀式、宇宙觀、傳統醫療與環境知識。

回答原則：
1. 先簡短摘要重點（3–5 點條列）。
2. 儘量說明「族名、場域、情境」與「知識來源背景」，避免抽象空話。
3. 若是推論或類比，要清楚標註「推測」而不是說成唯一正解。
4. 若資料不足或超出目前教材範圍，請誠實說明，並給出安全的延伸建議。
5. 全程使用繁體中文。"""

# ========= 載入模型 =========

print(f"載入模型：{MODEL_ID} ...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    use_auth_token=HF_TOKEN,
)

# 保險起見，若沒有 pad_token 就沿用 eos_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto",          # 自動分配到 GPU
    use_auth_token=HF_TOKEN,
)

model.eval()
print("模型載入完成。")


# ========= 建立提示詞 =========

def build_prompt(system_prompt: str, history: list[tuple[str, str]], user_message: str) -> str:
    """
    將 system_prompt + 歷史對話 + 新問題 組成一段文字 prompt。
    這裡用簡單的「使用者 / 助手」格式，對傳統知識生成已經很足夠。
    """
    system_prompt = system_prompt.strip()
    prompt = f"[系統提示]\n{system_prompt}\n\n"

    # 過去對話（若有）
    if history:
        prompt += "[對話紀錄]\n"
        for i, (user, bot) in enumerate(history, start=1):
            prompt += f"輪次 {i}：\n使用者：{user}\n助手：{bot}\n\n"

    # 最新一輪問題
    prompt += "[目前問題]\n"
    prompt += f"使用者：{user_message}\n助手："

    return prompt


# ========= 生成函式 =========

def generate_reply(
    user_message: str,
    chat_history: list[tuple[str, str]],
    system_prompt: str,
    temperature: float,
    max_new_tokens: int,
):
    if not user_message.strip():
        return chat_history, gr.update(value="")

    # 組合成一個大 prompt
    prompt_text = build_prompt(system_prompt, chat_history, user_message)

    inputs = tokenizer(
        prompt_text,
        return_tensors="pt",
        add_special_tokens=True,
    ).to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=True,
            temperature=float(temperature),
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id,
        )

    # 只取新生成的部分
    input_len = inputs["input_ids"].shape[-1]
    generated_tokens = outputs[0, input_len:]
    answer = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()

    chat_history = chat_history + [(user_message, answer)]
    return chat_history, ""  # 清空輸入框


def clear_history():
    return [], ""


# ========= Gradio 介面 =========

with gr.Blocks(title="語言橋傳統知識聊天機器人 — Mistral7B TK", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 語言橋傳統知識聊天機器人 — Mistral7B TK
        使用你自訓練的 **mistral7b-tk-sft-20251019-merged** 模型，離線在 Hugging Face Space 上回答與傳統知識相關的問題。

        建議題材舉例：
        - 不同族群對「夢境」的五種層次與詮釋差異  
        - 布農族狩獵儀式與祖靈信仰  
        - 排灣族階級制度與紋面、圖騰的意義  
        - 阿美族年齡階層制與植物分類知識  
        - 海外原住民（如 Inuit）對身體、疾病與療癒的理解  
        """
    )

    with gr.Row():
        # 左側：設定區
        with gr.Column(scale=1):
            system_prompt_box = gr.Textbox(
                label="系統提示（模型角色與回答風格）",
                value=DEFAULT_SYSTEM_PROMPT,
                lines=16,
            )
            temperature_slider = gr.Slider(
                label="溫度（創造性）",
                minimum=0.1,
                maximum=1.5,
                value=0.7,
                step=0.05,
            )
            max_tokens_slider = gr.Slider(
                label="最大回覆長度（token 數，大約字數的 1.5–2 倍）",
                minimum=64,
                maximum=1024,
                value=512,
                step=16,
            )
            gr.Markdown(
                """
                **小提醒：**
                - 回答太發散 → 降低溫度（0.4–0.7）。  
                - 回答太短 → 拉高「最大回覆長度」。  
                - Space 若常 timeout，可以稍微降低最大回覆長度。
                """
            )

        # 右側：聊天區
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                label="傳統知識 Chatbot",
                height=480,
                show_copy_button=True,
            )
            user_input = gr.Textbox(
                label="輸入你的問題（可多輪對話）",
                placeholder="例如：請比較布農族、排灣族和阿美族對治療疾病與夢境預兆的不同理解方式。",
                lines=4,
            )
            with gr.Row():
                send_btn = gr.Button("送出問題", variant="primary")
                clear_btn = gr.Button("清除對話")

    # 狀態：對話歷史
    state = gr.State([])  # list[tuple[user, bot]]

    # 綁定互動
    send_btn.click(
        fn=generate_reply,
        inputs=[
            user_input,
            state,
            system_prompt_box,
            temperature_slider,
            max_tokens_slider,
        ],
        outputs=[chatbot, user_input],
    ).then(
        fn=lambda h: h,
        inputs=[chatbot],
        outputs=[state],
    )

    user_input.submit(
        fn=generate_reply,
        inputs=[
            user_input,
            state,
            system_prompt_box,
            temperature_slider,
            max_tokens_slider,
        ],
        outputs=[chatbot, user_input],
    ).then(
        fn=lambda h: h,
        inputs=[chatbot],
        outputs=[state],
    )

    clear_btn.click(
        fn=clear_history,
        inputs=[],
        outputs=[chatbot, user_input],
    ).then(
        fn=lambda: [],
        inputs=[],
        outputs=[state],
    )

# 在 HF Space 中不需要 demo.launch()，平台會自動呼叫 demo