import re
import json
import requests
import os
import gradio as gr

HF_TOKEN = os.getenv("HF_TOKEN")

task_options = ["重點整理", "問答", "翻譯"]

MODEL_URLS = {
    "重點整理": {
        "翻譯": "https://api-inference.huggingface.co/models/Helsinki-NLP/opus-mt-en-zh",
        "摘要": "https://api-inference.huggingface.co/models/csebuetnlp/mT5_multilingual_XLSum"
    },
    "問答": "https://api-inference.huggingface.co/models/luhua/chinese_pretrain_mrc_roberta_wwm_ext_large",
    "翻譯": "https://api-inference.huggingface.co/models/Helsinki-NLP/opus-mt-en-zh"
}

def clean_rtf(text):
    text = re.sub(r"\'.", "", text)
    text = re.sub(r"\[a-z]+[0-9]* ?", "", text)
    text = re.sub(r"[{}]", "", text)
    return text.strip()

def run_model(task, text, question=None):
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type": "application/json"
    }

    if task == "重點整理":
        translation_payload = {"inputs": text}
        translation_url = MODEL_URLS["重點整理"]["翻譯"]
        translation_response = requests.post(translation_url, headers=headers, json=translation_payload)

        if translation_response.status_code != 200:
            return "❌ 翻譯失敗：" + translation_response.text

        try:
            translated_text = translation_response.json()[0]['translation_text']
        except Exception as e:
            return f"❌ 翻譯結果解析錯誤：{str(e)}"

        summarization_payload = {"inputs": f"summarize: {translated_text}"}
        summarization_url = MODEL_URLS["重點整理"]["摘要"]
        summarization_response = requests.post(summarization_url, headers=headers, json=summarization_payload)

        if summarization_response.status_code != 200:
            return "❌ 摘要失敗：" + summarization_response.text

        try:
            summary_text = summarization_response.json()[0]['summary_text']
        except Exception as e:
            return f"❌ 摘要結果解析錯誤：{str(e)}"

        return summary_text

    elif task == "問答":
        if not question:
            return "❌ 請輸入問題"

        qa_payload = {
            "inputs": {
                "question": question,
                "context": text
            }
        }
        qa_url = MODEL_URLS["問答"]
        qa_response = requests.post(qa_url, headers=headers, json=qa_payload)

        if qa_response.status_code != 200:
            return "❌ 問答失敗：" + qa_response.text

        try:
            answer = qa_response.json().get("answer", "⚠️ 找不到答案")
            trans_payload = {"inputs": answer}
            trans_url = MODEL_URLS["翻譯"]
            trans_response = requests.post(trans_url, headers=headers, json=trans_payload)
            if trans_response.status_code == 200:
                return trans_response.json()[0]['translation_text']
            else:
                return answer
        except Exception as e:
            return f"❌ 回答解析錯誤：{str(e)}"

    elif task == "翻譯":
        translation_payload = {"inputs": text}
        translation_url = MODEL_URLS["翻譯"]
        translation_response = requests.post(translation_url, headers=headers, json=translation_payload)

        if translation_response.status_code != 200:
            return "❌ 翻譯失敗：" + translation_response.text

        try:
            return translation_response.json()[0]['translation_text']
        except Exception as e:
            return f"❌ 翻譯結果解析錯誤：{str(e)}"

    else:
        return "❌ 不支援的任務"

with gr.Blocks() as demo:
    gr.Markdown("# 🌐 多功能語言處理器（繁體中文）\n支援：重點整理（英文→中）、問答、翻譯（英翻中）")

    with gr.Row():
        task = gr.Dropdown(choices=task_options, label="請選擇任務")

    with gr.Row():
        text_input = gr.Textbox(lines=10, label="輸入文章 / 內容")

    with gr.Row():
        question_input = gr.Textbox(label="問題（問答任務用）", placeholder="選擇問答任務時必填")

    with gr.Row():
        output = gr.Textbox(lines=5, label="輸出結果")

    with gr.Row():
        run_button = gr.Button("執行")

    run_button.click(fn=run_model, inputs=[task, text_input, question_input], outputs=output)

if __name__ == "__main__":
    demo.launch()