# app.py —— 永久一键切换：微调 / 测试 / 机器人
import os
import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer, TextDataset, DataCollatorForLanguageModeling
from transformers import Trainer, TrainingArguments
import re
import torch

MODEL_DIR = "./fine_tuned_model"
BASE_MODEL = "uer/gpt2-chinese-cluecorpussmall"

# ==================== 微调函数 ====================
def finetune_model(progress=gr.Progress()):
    if os.path.exists(MODEL_DIR):
        return "已存在微调好的模型，无需重复微调！直接切换到「测试」或「机器人」模式即可。"

    progress(0, desc="正在读取家谱文件...")
    with open("genealogy.txt", encoding="utf-8") as f:
        text = f.read()

    # 简单但有效的信息抽取
    blocks = re.split(r"\n\d+——|\n故祖考|\n故祖妣|\n孝男|\n孝媳|\n孝孙|\n曾孙|\n外曾孙|\n外玄孙", text)
    examples = []
    for block in blocks[1:]:
        if len(block) < 10:
            continue
        # 提取姓名
        name_match = re.search(r"[\u4e00-\u9fa5]{2,4}", block[:15])
        name = name_match.group() if name_match else "某人"

        born = re.search(r"生[於于]\s*([^\n地名]+)", block)
        died = re.search(r"(死[於于]|卒于|去世于)\s*([^\n地名]+)", block)
        place = re.search(r"地名\s*([^\n]+)", block)

        born = born.group(1).strip() if born else "未知"
        place = place.group(1).strip() if place else "未知"
        died = died.group(2).strip() if died else None

        examples.append(f"查询: {name}出生时间和地点？\n回答: {name}出生于{born}，地点{place}。")
        if died:
            examples.append(f"查询: {name}去世时间？\n回答: {name}去世于{died}。")

    data_text = "\n\n".join(examples)
    with open("train.txt", "w", encoding="utf-8") as f:
        f.write(data_text)

    progress(0.3, desc="加载 tokenizer 和数据集...")
    tokenizer = GPT2Tokenizer.from_pretrained(BASE_MODEL)
    dataset = TextDataset(tokenizer=tokenizer, file_path="train.txt", block_size=128)
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    progress(0.6, desc="开始微调（预计40–90分钟，请耐心等待勿关闭页面）...")
    model = GPT2LMHeadModel.from_pretrained(BASE_MODEL)

    training_args = TrainingArguments(
        output_dir="./results",
        overwrite_output_dir=True,
        num_train_epochs=3,
        per_device_train_batch_size=2,
        save_steps=500,
        logging_steps=20,
        save_total_limit=1,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=dataset,
    )
    trainer.train()

    progress(1.0, desc="保存模型...")
    model.save_pretrained(MODEL_DIR)
    tokenizer.save_pretrained(MODEL_DIR)

    return "微调成功！模型已保存。现在可以切换到「测试」或「机器人」模式使用啦！"

# ==================== 推理函数 ====================
def load_model():
    if not os.path.exists(MODEL_DIR):
        return None, None
    tokenizer = GPT2Tokenizer.from_pretrained(MODEL_DIR)
    model = GPT2LMHeadModel.from_pretrained(MODEL_DIR)
    return model, tokenizer

model, tokenizer = load_model()

def chat(message, history):
    if model is None:
        return "模型尚未微调完成！请先在上面选择「微调模型」运行一次。"
    
    prompt = "\n".join([f"查询: {h[0]}\n回答: {h[1]}" for h in history if h[1]] + [f"查询: {message}\n回答:"])
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs,
        max_length=inputs.shape[1] + 120,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        no_repeat_ngram_size=3,
        pad_token_id=tokenizer.eos_token_id
    )
    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return reply.split("回答:")[-1].strip()

# ==================== 主界面 ====================
with gr.Blocks(theme=gr.themes.Soft(), title="冯氏家谱智能查询") as demo:
    gr.Markdown("# 冯氏家谱智能查询机器人\n一次上传，永久使用，三个模式一键切换")

    mode = gr.Radio(
        ["微调模型", "测试模型", "正式机器人"],
        label="请选择当前模式",
        value="正式机器人" if os.path.exists(MODEL_DIR) else "微调模型"
    )

    with gr.Column(visible=False) as finetune_box:
        gr.Markdown("### 第一步：点击下方按钮开始微调（只需执行一次，耐心等待40–90分钟）")
        btn = gr.Button("开始微调模型（免费CPU可完成）", variant="primary", size="lg")
        output = gr.Textbox(label="微调日志", lines=20)
        btn.click(finetune_model, None, output)

    with gr.Column(visible=True) as chat_box:
        gr.ChatInterface(
            fn=chat,
            title=None,
            description="直接问：冯达尊出生地、冯乔福几个儿子、冯宗德姐妹是谁……",
            examples=[
                "冯达尊出生时间和地点？",
                "冯乔源的孩子有哪些？",
                "冯宗福的兄弟姐妹是谁？",
                "冯永明是哪一支的？",
                "冯仁杰的父亲是谁？"
            ]
        )

    def update_interface(choice):
        if choice == "微调模型":
            return gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=False), gr.update(visible=True)

    mode.change(update_interface, mode, [finetune_box, chat_box])

# 启动时自动显示正确界面
demo.launch()