"""
Inference module -- load checkpoint tu HuggingFace va generate SQL tu model
input string. Ho tro day du catalog checkpoint cua du an (4 stage x 3
architecture = 12 model), nguoi dung TU CHON qua dropdown UI bang 1 model_id
duy nhat (xem MODEL_CATALOG) -- khong con phu thuoc cung mode du lieu dang
active (fixed/uploaded), vi cac checkpoint Stage 2/3 da verify hoat dong tot
tren CA HAI dang schema (Spider-style co PK/FK VA WikiSQL-style khong PK/FK,
xem ghi chu duoi).

model_id format: "{terminal_dataset}_stage{n}_{architecture}"
  - terminal_dataset: dataset CUOI CUNG model duoc fine-tune tren (quyet dinh
    quy uoc sinh SQL ma model da hoc, KHONG can trung voi schema dang dung).
  - stage: 1 (train rieng le) | 2 (WikiSQL->Spider) | 3 (WikiSQL->Spider->BIRD)

Da verify thuc nghiem (2026-06-26): checkpoint Stage 2/3 (da tiep xuc Spider/
BIRD multi-column SELECT + GROUP BY) VAN sinh dung multi-column SELECT ngay ca
khi duoc dua schema dang WikiSQL (khong PK/FK, vd CSV upload) -- vd
"SELECT store_name, revenue FROM data_table". Day la ly do chinh de mo rong
catalog nay: checkpoint Stage 1 WikiSQL (terminal_dataset=wikisql) KHONG BAO
GIO sinh duoc 2 cot (gioi han cau truc cua chinh grammar WikiSQL, xem
visualizer.py/CLAUDE.md) -> khong the ra chart; Stage 2/3 thi co the, doi lai
do chinh xac cho cau don gian (SELECT 1 cot, WHERE) co the thap hon mot chut
so voi Stage 1 WikiSQL thuan (xem EM/EX tung checkpoint trong label).

Luu y dinh dang tokenizer KHONG nhat quan giua cac checkpoint trong du an nay
(mot so luu san tokenizer.json FAST, mot so chi co vocab.json/merges.txt SLOW)
-- PHAI tu phat hien dinh dang co san (_has_fast_tokenizer_file) thay vi
hardcode use_fast=True/False.
"""

import os

import torch
from huggingface_hub import snapshot_download
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

NUM_BEAMS = 4

ARCHITECTURE_NAMES = {"t5": "T5-base", "codet5": "CodeT5-base", "bart": "BART-base"}

# Moi entry: repo, ckpt, max_input/output_length, label (hien thi dropdown),
# stage (1/2/3), terminal_dataset ("wikisql"|"spider"|"bird" -- dataset CUOI
# CUNG model fine-tune tren, quyet dinh quy uoc quote khi post-process).
MODEL_CATALOG = {
    # ---- Stage 1: train rieng le tren WikiSQL (EM greedy, in-training) ----
    "wikisql_stage1_t5": {
        "repo": "minimew/text2sql-checkpoints",
        "ckpt": "wikisql/t5-base/run_6/epoch-12-em0_6918",
        "max_input_length": 384, "max_output_length": 64,
        "architecture": "t5", "stage": 1, "terminal_dataset": "wikisql",
        "label": "T5-base · Stage 1 (WikiSQL)",
    },
    "wikisql_stage1_codet5": {
        "repo": "minimew/text2sql-checkpoint-2",
        "ckpt": "wikisql/codet5-base/run_1/epoch-08-em0_7688",
        "max_input_length": 384, "max_output_length": 64,
        "architecture": "codet5", "stage": 1, "terminal_dataset": "wikisql",
        "label": "CodeT5-base · Stage 1 (WikiSQL)",
    },
    "wikisql_stage1_bart": {
        "repo": "minimew/text2sql-checkpoints",
        "ckpt": "wikisql/bart-base/run_1/epoch-05-em0_7228",
        "max_input_length": 384, "max_output_length": 64,
        "architecture": "bart", "stage": 1, "terminal_dataset": "wikisql",
        "label": "BART-base · Stage 1 (WikiSQL)",
    },
    # ---- Stage 1: train rieng le tren Spider (EM/EX final eval, beam=4) ----
    "spider_stage1_t5": {
        "repo": "minimew/text2sql-checkpoints",
        "ckpt": "spider/t5-base/run_5/epoch-07-em0_3956",
        "max_input_length": 1024, "max_output_length": 128,
        "architecture": "t5", "stage": 1, "terminal_dataset": "spider",
        "label": "T5-base · Stage 1 (Spider)",
    },
    "spider_stage1_codet5": {
        "repo": "minimew/text2sql-checkpoint-2",
        "ckpt": "spider/codet5-base/run_1/epoch-07-em0_4391",
        "max_input_length": 1024, "max_output_length": 128,
        "architecture": "codet5", "stage": 1, "terminal_dataset": "spider",
        "label": "CodeT5-base · Stage 1 (Spider)",
    },
    "spider_stage1_bart": {
        "repo": "minimew/text2sql-checkpoints",
        "ckpt": "spider/bart-base/run_1/epoch-08-em0_2940",
        "max_input_length": 1024, "max_output_length": 128,
        "architecture": "bart", "stage": 1, "terminal_dataset": "spider",
        "label": "BART-base · Stage 1 (Spider)",
    },
    # ---- Stage 2: WikiSQL -> Spider (EM/EX final eval, beam=4) ----
    # Da verify: van sinh dung multi-column SELECT tren schema kieu WikiSQL
    # (CSV upload) -- lua chon tot nhat neu can chart tu CSV.
    "spider_stage2_t5": {
        "repo": "minimew/text2sql-checkpoints",
        "ckpt": "stage2/spider/t5-base/run_1/epoch-08-em0_3752",
        "max_input_length": 1024, "max_output_length": 128,
        "architecture": "t5", "stage": 2, "terminal_dataset": "spider",
        "label": "T5-base · Stage 2 (WikiSQL→Spider)",
    },
    "spider_stage2_codet5": {
        "repo": "minimew/text2sql-checkpoint-2",
        "ckpt": "stage2/spider/codet5-base/run_1/epoch-07-em0_4275",
        "max_input_length": 1024, "max_output_length": 128,
        "architecture": "codet5", "stage": 2, "terminal_dataset": "spider",
        "label": "CodeT5-base · Stage 2 (WikiSQL→Spider)",
    },
    "spider_stage2_bart": {
        "repo": "minimew/text2sql-checkpoints",
        "ckpt": "stage2/spider/bart-base/run_1/epoch-07-em0_3037",
        "max_input_length": 1024, "max_output_length": 128,
        "architecture": "bart", "stage": 2, "terminal_dataset": "spider",
        "label": "BART-base · Stage 2 (WikiSQL→Spider)",
    },
    # ---- Stage 3: WikiSQL -> Spider -> BIRD (EX tren BIRD dev, final eval) ----
    # Da verify: van sinh SQL hop le tren CA schema Spider va WikiSQL (khong
    # can field "evidence:" -- demo nay khong co BIRD DB thuc/evidence injection).
    "bird_stage3_t5": {
        "repo": "minimew/text2sql-checkpoints",
        "ckpt": "stage3/bird/t5-base/run_1/epoch-12-em0_0495",
        "max_input_length": 1024, "max_output_length": 128,
        "architecture": "t5", "stage": 3, "terminal_dataset": "bird",
        "label": "T5-base · Stage 3 (→BIRD)",
    },
    "bird_stage3_codet5": {
        "repo": "minimew/text2sql-checkpoint-2",
        "ckpt": "stage3/bird/codet5-base/run_1/epoch-11-em0_1037",
        "max_input_length": 1024, "max_output_length": 128,
        "architecture": "codet5", "stage": 3, "terminal_dataset": "bird",
        "label": "CodeT5-base · Stage 3 (→BIRD)",
    },
    "bird_stage3_bart": {
        "repo": "minimew/text2sql-checkpoints",
        "ckpt": "stage3/bird/bart-base/run_1/epoch-12-em0_0443",
        "max_input_length": 1024, "max_output_length": 128,
        "architecture": "bart", "stage": 3, "terminal_dataset": "bird",
        "label": "BART-base · Stage 3 (→BIRD)",
    },
    # ---- Joint: train dong thoi WikiSQL + Spider + BIRD (40/40/20 mix) ----
    # CodeT5 only (RQ5). Spider EX 60.74% / WikiSQL EX 81.89% / BIRD EX 23.08%.
    # terminal_dataset="spider": 60% du lieu la Spider+BIRD dung single-quote
    # cho string value -- ap dung quote normalization nhu cac checkpoint Spider.
    "joint_codet5": {
        "repo": "minimew/text2sql-checkpoint-2",
        "ckpt": "joint/codet5-base/run_1/epoch-12-emavg0_4333",
        "max_input_length": 1024, "max_output_length": 128,
        "architecture": "codet5", "stage": "joint", "terminal_dataset": "spider",
        "label": "CodeT5-base · Joint (WikiSQL + Spider + BIRD)",
    },
}

DEFAULT_MODEL_ID_FOR_MODE = {
    "fixed": "spider_stage1_codet5",
    "uploaded": "wikisql_stage1_codet5",
}

_state = {}


def _has_fast_tokenizer_file(ckpt_path: str) -> bool:
    return os.path.exists(os.path.join(ckpt_path, "tokenizer.json"))


def get_checkpoint_config(model_id: str) -> dict:
    return MODEL_CATALOG[model_id]


def load_model(model_id: str):
    """Load model + tokenizer cho 1 checkpoint, cache rieng theo model_id.
    Goi lai voi cung model_id se tra ve doi tuong da cache."""
    if model_id in _state:
        return _state[model_id]["model"], _state[model_id]["tokenizer"], _state[model_id]["device"]

    cfg = get_checkpoint_config(model_id)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # local_dir_use_symlinks=False: tranh OSError WinError 1314 tren Windows
    # khi Developer Mode chua kich hoat (huggingface_hub mac dinh dung symlink
    # de cache, can quyen dac biet de tao symlink tren Windows).
    local_dir = snapshot_download(
        repo_id=cfg["repo"], allow_patterns=[f"{cfg['ckpt']}/*"], local_dir_use_symlinks=False
    )
    ckpt_path = os.path.join(local_dir, cfg["ckpt"])

    use_fast = _has_fast_tokenizer_file(ckpt_path)
    tokenizer = AutoTokenizer.from_pretrained(ckpt_path, use_fast=use_fast)
    model = AutoModelForSeq2SeqLM.from_pretrained(ckpt_path).to(device)
    model.eval()

    _state[model_id] = {"model": model, "tokenizer": tokenizer, "device": device}
    return model, tokenizer, device


def generate_sql(model_input: str, model_id: str) -> str:
    """model_input: full string 'question: ... | schema: ...'."""
    model, tokenizer, device = load_model(model_id)
    cfg = get_checkpoint_config(model_id)

    inputs = tokenizer(
        model_input,
        return_tensors="pt",
        truncation=True,
        max_length=cfg["max_input_length"],
    ).to(device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            num_beams=NUM_BEAMS,
            max_length=cfg["max_output_length"],
            early_stopping=True,
        )

    sql = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

    # Dong bo voi normalize convention da dung luc eval CHO SPIDER/BIRD: model
    # doi khi sinh double-quote cho STRING VALUE ("France") nhung gold dung
    # single-quote ('France') -- doi sang single-quote de tang ty le execute
    # thanh cong. KHONG ap dung cho checkpoint terminal_dataset=wikisql: o
    # format WikiSQL, double-quote dung de quote TEN COT (xem quote_column()
    # trong process_wikisql.py), con string value đa dung single-quote san --
    # thay the se bien ten cot co quote thanh string literal sai (vd
    # "product name" -> 'product name' bi SQLite hieu thanh chuoi hang so
    # thay vi tham chieu cot).
    if cfg["terminal_dataset"] != "wikisql":
        sql = sql.replace('"', "'")
    return sql