import os
import shutil
import zipfile

import gradio as gr
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer


ZIP_PATH = "model.zip"
EXTRACT_DIR = "./model"
MODEL_DIR = os.path.join(EXTRACT_DIR, "final_model")
DEVICE = torch.device("cpu")


def prepare_model_files():
    config_path = os.path.join(MODEL_DIR, "config.json")
    weights_path = os.path.join(MODEL_DIR, "model.safetensors")

    if os.path.isfile(config_path) and os.path.isfile(weights_path):
        print("模型已經解壓完成。", flush=True)
        return

    if not os.path.isfile(ZIP_PATH):
        raise FileNotFoundError("找不到 model.zip")

    if os.path.exists(EXTRACT_DIR):
        shutil.rmtree(EXTRACT_DIR)

    os.makedirs(EXTRACT_DIR, exist_ok=True)

    print("正在解壓模型……", flush=True)

    with zipfile.ZipFile(ZIP_PATH, "r") as zip_ref:
        zip_ref.extractall(EXTRACT_DIR)

    print("模型解壓完成。", flush=True)

    if not os.path.isfile(config_path):
        raise FileNotFoundError(
            f"找不到 {config_path}，請確認 ZIP 裡面有 final_model 資料夾。"
        )

    if not os.path.isfile(weights_path):
        raise FileNotFoundError(
            f"找不到 {weights_path}，請確認 model.safetensors 已包含在 ZIP 中。"
        )


prepare_model_files()

print("正在載入 tokenizer……", flush=True)

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_DIR,
    local_files_only=True
)

print("正在載入 BERT 模型到 CPU……", flush=True)

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_DIR,
    local_files_only=True
)

model.to(DEVICE)
model.eval()

print("模型載入成功。", flush=True)


LABEL_MAP = {
    0: "詐騙訊息",
    1: "真實訊息"
}


def predict(text: str) -> str:
    text = text.strip()

    if not text:
        return "請先輸入要判斷的訊息。"

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=256
    )

    inputs = {
        key: value.to(DEVICE)
        for key, value in inputs.items()
    }

    with torch.inference_mode():
        outputs = model(**inputs)
        probabilities = torch.softmax(outputs.logits, dim=-1)[0]

    probabilities = probabilities.cpu()

    pred_id = int(torch.argmax(probabilities).item())
    confidence = float(probabilities[pred_id].item() * 100)

    scam_probability = float(probabilities[0].item() * 100)
    real_probability = float(probabilities[1].item() * 100)

    pred_label = LABEL_MAP.get(
        pred_id,
        f"未知類別 {pred_id}"
    )

    if pred_id == 0:
        advice = (
            "請勿立即匯款、提供銀行帳號、密碼、信用卡資料或簡訊驗證碼。"
            "建議透過官方管道查證，必要時撥打 165 反詐騙專線。"
        )
    else:
        advice = (
            "模型判斷較接近真實訊息，但模型仍可能誤判。"
            "若內容涉及金錢、帳戶或個人資料，仍應再次查證。"
        )

    return (
        f"🔍 判斷結果：{pred_label}\n\n"
        f"📊 模型信心度：{confidence:.2f}%\n\n"
        f"🚨 詐騙機率：{scam_probability:.2f}%\n"
        f"✅ 真實機率：{real_probability:.2f}%\n\n"
        f"⚠️ 提醒：{advice}"
    )


demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(
        lines=6,
        label="訊息內容",
        placeholder="請輸入要判斷的新聞、貼文或可疑訊息……"
    ),
    outputs=gr.Textbox(
        label="辨識結果",
        lines=8
    ),
    title="LINE Bot BERT 詐騙／真實訊息辨識 API",
    description=(
        "輸入文字後，系統會判斷內容較接近詐騙訊息或真實訊息。"
        "建議輸入較完整的內容，文字過短可能降低辨識效果。"
    ),
    api_name="predict"
)


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False,
        show_error=True
    )