import json
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file

from model import ModernBertConfig, ModernBertModel


class PieceCharTokenizer:
    def __init__(self, model_dir):
        import piece_tokenizer as pt
        model_dir = Path(model_dir)
        self._tok = pt.Tokenizer()
        self._tok.load(str(model_dir / "piece.model"), cn_dict="no")
        mask_path = model_dir / "mask_token_id.txt"
        self.mask_token_id = int(mask_path.read_text().strip()) if mask_path.exists() else self._tok.vocab_size()
        self.vocab_size = self._tok.vocab_size() + 1
        self.pad_token_id = self._tok.piece_to_id("<pad>")
        self.unk_token_id = 0
        self.cache = {}
        self.inv_cache = {}

    def char_to_id(self, char):
        if char in self.cache:
            return self.cache[char]
        ids = self._tok.encode_as_ids(char)
        tid = ids[0] if ids else self.unk_token_id
        self.cache[char] = tid
        self.inv_cache.setdefault(tid, char)
        return tid


class BERTcForCSC(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.bert = ModernBertModel(config)
        self.cor_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.cor_head.weight = self.bert.embed.weight
        self.det_head = nn.Linear(config.hidden_size, 1)

    @classmethod
    def from_pretrained(cls, model_dir, map_location="cpu"):
        model_dir = Path(model_dir)
        cfg = ModernBertConfig(**json.loads((model_dir / "config.json").read_text()))
        model = cls(cfg)
        state = load_file(str(model_dir / "model.safetensors"), device=str(map_location))
        missing, unexpected = model.load_state_dict(state, strict=False)
        allowed_missing = {"cor_head.weight"}
        if set(missing) != allowed_missing or unexpected:
            raise RuntimeError(f"Bad state dict: missing={missing}, unexpected={unexpected}")
        model.cor_head.weight = model.bert.embed.weight
        model.eval()
        return model

    def forward(self, input_ids, attention_mask=None):
        h = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        cor_logits = self.cor_head(h)
        det_logits = self.det_head(h).squeeze(-1)
        return cor_logits, det_logits

    @torch.no_grad()
    def correct(self, texts, tokenizer, threshold=0.7, max_len=128, device=None):
        if isinstance(texts, str):
            single = True
            texts = [texts]
        else:
            single = False
        device = device or next(self.parameters()).device
        self.eval()
        lengths = [min(len(t), max_len) for t in texts]
        max_l = max(lengths) if lengths else 0
        input_ids = torch.full((len(texts), max_l), tokenizer.pad_token_id, dtype=torch.long, device=device)
        attn = torch.zeros((len(texts), max_l), dtype=torch.long, device=device)
        for i, text in enumerate(texts):
            ids = [tokenizer.char_to_id(c) for c in text[:lengths[i]]]
            if ids:
                input_ids[i, :len(ids)] = torch.tensor(ids, dtype=torch.long, device=device)
                attn[i, :len(ids)] = 1
        cor_logits, _ = self(input_ids, attn)
        probs = F.softmax(cor_logits, dim=-1)
        top_probs, top_ids = probs.max(dim=-1)
        out = []
        for i, text in enumerate(texts):
            chars = list(text[:lengths[i]])
            pred = []
            for j, orig in enumerate(chars):
                tid = int(top_ids[i, j].item())
                prob = float(top_probs[i, j].item())
                pred.append(tokenizer.inv_cache.get(tid, orig) if prob >= threshold else orig)
            if len(text) > lengths[i]:
                pred.extend(list(text[lengths[i]:]))
            out.append("".join(pred))
        return out[0] if single else out