"""Inference cho Vietnamese QA Stacking Ensemble v1. Load từ Hugging Face Hub."""
import json
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForQuestionAnswering, AutoTokenizer

try:
    from huggingface_hub import hf_hub_download
except ImportError:
    hf_hub_download = None


class MetaCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv1d(4, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv1d(32, 32, 3, padding=1),
            nn.ReLU(),
        )
        self.start_fc = nn.Linear(32, 1)
        self.end_fc = nn.Linear(32, 1)

    def forward(self, x):
        x = self.conv(x)
        start = self.start_fc(x.transpose(1, 2)).squeeze(-1)
        end = self.end_fc(x.transpose(1, 2)).squeeze(-1)
        return start, end


def _load_tok(mid, use_fast=True):
    try:
        return AutoTokenizer.from_pretrained(mid, use_fast=use_fast)
    except Exception as e:
        if "sentencepiece" in str(e).lower() and use_fast:
            return AutoTokenizer.from_pretrained(mid, use_fast=False)
        raise


def load_ensemble(repo_id: str = None, local_dir: str = None):
    """Load ensemble từ Hugging Face hoặc thư mục local."""
    if local_dir:
        path = Path(local_dir)
        config_path = path / "config.json"
        meta_path = path / "meta_cnn.pth"
    elif repo_id and hf_hub_download:
        config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
        meta_path = hf_hub_download(repo_id=repo_id, filename="meta_cnn.pth")
    else:
        raise ValueError("Cần repo_id hoặc local_dir")

    with open(config_path, encoding="utf-8") as f:
        config = json.load(f)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    tokenizer1 = _load_tok(config["base_models"][0])
    tokenizer2 = AutoTokenizer.from_pretrained(config["base_models"][1], use_fast=False)

    model1 = AutoModelForQuestionAnswering.from_pretrained(config["base_models"][0]).to(device)
    model2 = AutoModelForQuestionAnswering.from_pretrained(config["base_models"][1]).to(device)

    meta_model = MetaCNN().to(device)
    meta_model.load_state_dict(torch.load(meta_path, map_location=device))
    meta_model.eval()
    model1.eval()
    model2.eval()

    return {
        "tokenizer1": tokenizer1,
        "tokenizer2": tokenizer2,
        "model1": model1,
        "model2": model2,
        "meta_model": meta_model,
        "device": device,
        "max_len_1": config.get("max_length", 512),
        "max_len_2": config.get("max_len_2", 256),
    }


def _pad_to_512(x):
    if x.size(0) < 512:
        pad = torch.zeros(512 - x.size(0), dtype=x.dtype, device=x.device)
        x = torch.cat([x, pad], dim=0)
    return x[:512]


def predict(question: str, context: str, ensemble: dict, max_answer_len: int = 30):
    """Trả về (answer, no_answer_probability)."""
    t1, t2 = ensemble["tokenizer1"], ensemble["tokenizer2"]
    m1, m2 = ensemble["model1"], ensemble["model2"]
    meta = ensemble["meta_model"]
    dev = ensemble["device"]
    max1, max2 = ensemble["max_len_1"], ensemble["max_len_2"]

    enc1 = t1(
        question,
        context,
        return_tensors="pt",
        truncation="only_second",
        max_length=max1,
        padding="max_length",
    )
    enc2 = t2(
        question,
        context,
        return_tensors="pt",
        truncation="only_second",
        max_length=max2,
        padding="max_length",
    )

    inp1 = {k: v.to(dev) for k, v in enc1.items()}
    inp2 = {k: v.to(dev) for k, v in enc2.items()}

    try:
        seq_ids = enc1.sequence_ids(0)
    except Exception:
        sep_id = t1.convert_tokens_to_ids(t1.sep_token or "</s>")
        ids = enc1["input_ids"][0].tolist()
        sep_pos = [i for i, x in enumerate(ids) if x == sep_id]
        if len(sep_pos) < 2:
            return "", 1.0
        ctx_idx = list(range(sep_pos[0] + 1, sep_pos[1]))
    else:
        ctx_idx = [i for i, s in enumerate(seq_ids) if s == 1]

    if not ctx_idx:
        return "", 1.0

    ctx_start, ctx_end = ctx_idx[0], ctx_idx[-1]

    with torch.no_grad():
        o1, o2 = m1(**inp1), m2(**inp2)

    s1 = o1.start_logits[0][:512]
    e1 = o1.end_logits[0][:512]
    s2 = _pad_to_512(o2.start_logits[0])
    e2 = _pad_to_512(o2.end_logits[0])

    combined = torch.stack([s1, e1, s2, e2], dim=0).unsqueeze(0)
    with torch.no_grad():
        fs, fe = meta(combined)

    sp = F.softmax(fs[0], dim=-1)
    ep = F.softmax(fe[0], dim=-1)

    best = -1e9
    bs, be = ctx_start, ctx_start
    for s in range(ctx_start, ctx_end + 1):
        for e in range(s, min(s + max_answer_len, ctx_end) + 1):
            sc = torch.log(sp[s] + 1e-12) + torch.log(ep[e] + 1e-12)
            if sc > best:
                best, bs, be = sc, s, e

    null = torch.log(sp[0] + 1e-12) + torch.log(ep[0] + 1e-12)
    no_ans = torch.sigmoid(null - best).item()
    if null > best:
        return "", no_ans

    ans = t1.decode(enc1["input_ids"][0][bs:be + 1], skip_special_tokens=True).strip()
    return ans, no_ans