import gradio as gr
import torch
import json
import html
import traceback
from transformers import AutoModelForCausalLM, AutoTokenizer

print("Loading model...")
model_name = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)
print("Model loaded.")

SYSTEM_PROMPT = """You are an English learning assistant. Extract 8-20 useful expressions from the text.
For each expression, output a JSON object with keys: expression, meaning, explanation, original_context, extra_example.
Meaning and explanation should be in Chinese.
Output must be a JSON array. No extra text."""

def analyze(text):
    try:
        if not text or len(text.strip()) < 20:
            return "<div style='color:red'>⚠️ Please enter at least 20 characters.</div>"

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": text}
        ]
        inputs = tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt"
        ).to(model.device)

        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_new_tokens=1024,
                do_sample=False,
                temperature=1.0
            )

        response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)

        # 提取 JSON
        if "```json" in response:
            response = response.split("```json")[1].split("```")[0]
        elif "```" in response:
            response = response.split("```")[1].split("```")[0]
        start = response.find("[")
        end = response.rfind("]") + 1
        if start == -1 or end == 0:
            return f"<div style='color:red'>No JSON array found. Raw response:<br>{html.escape(response[:300])}</div>"

        json_str = response[start:end]
        data = json.loads(json_str)

        cards = ""
        for e in data:
            cards += f"""
            <div style="background:white;border-radius:16px;border:1px solid #ddd;padding:1rem;margin-bottom:1rem;">
                <b style="font-size:1.2rem;">{html.escape(str(e.get('expression', '')))}</b><br>
                <b>Meaning</b><br>{html.escape(str(e.get('meaning', '')))}<br>
                <b>Explanation</b><br>{html.escape(str(e.get('explanation', '')))}<br>
                <b>Original Context</b><br>{html.escape(str(e.get('original_context', '')))}<br>
                <b>Extra Example</b><br>{html.escape(str(e.get('extra_example', '')))}
            </div>
            """
        return cards if cards else "<div>No expressions extracted.</div>"
    except Exception as e:
        error_html = f"<div style='color:red; background:#ffe0e0; padding:1rem; border-radius:8px;'>"
        error_html += f"<b>Error:</b> {html.escape(str(e))}<br><br>"
        error_html += f"<details><summary>Full traceback</summary><pre>{html.escape(traceback.format_exc())}</pre></details>"
        error_html += "</div>"
        return error_html

# 浅色主题
theme = gr.themes.Soft(
    primary_hue="neutral",
    secondary_hue="neutral",
    font=gr.themes.GoogleFont("Inter"),
).set(
    body_background_fill="#fafaf9",
    button_primary_background_fill="#1a1a1a",
    button_primary_text_color="white",
    block_background_fill="white",
)

with gr.Blocks(theme=theme, title="InContext") as demo:
    gr.Markdown("# InContext\n### Learn English Expressions Through Real Content")
    with gr.Row():
        txt = gr.Textbox(lines=10, placeholder="Paste English content here...", label="")
    btn = gr.Button("Analyze", variant="primary")
    out = gr.HTML()
    btn.click(analyze, txt, out)

demo.launch()