import json
from functools import lru_cache
from pathlib import Path

import torch
from huggingface_hub import hf_hub_download

from model import LemmaModel, Vocab


MODEL_REPO_ID = "usmannawaz/ocscomdemo"


@lru_cache(maxsize=1)
def load_registry(registry_path="models_registry.json"):
    path = Path(__file__).resolve().parent / registry_path

    with path.open(encoding="utf8") as file:
        return json.load(file)


@lru_cache(maxsize=3)
def download_model_files(model_id):
    registry = load_registry()

    if model_id not in registry:
        raise KeyError(f"Model ID not found: {model_id}")

    item = registry[model_id]

    config_path = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        repo_type="model",
        filename=item["config_file"],
    )

    vocab_path = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        repo_type="model",
        filename=item["vocab_file"],
    )

    weights_path = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        repo_type="model",
        filename=item["model_file"],
    )

    return config_path, vocab_path, weights_path


class OldSlavicLemmatizer:
    def __init__(self, model, vocab, config, device):
        self.model = model
        self.vocab = vocab
        self.config = config
        self.device = torch.device(device)
        self.sep_char = config.get("sep_char", "⟂")
        self.k_context = int(config.get("k_context", 2))
        self.max_gen_len = int(config.get("max_gen_len", 30))

    def make_source(self, form, left_context=None, right_context=None):
        left_context = left_context or []
        right_context = right_context or []

        left = " ".join(
            left_context[-self.k_context:]
        ).strip()

        right = " ".join(
            right_context[:self.k_context]
        ).strip()

        src_left = left + " " if left else ""
        src_right = " " + right if right else ""

        return (
            f"{src_left}"
            f"{self.sep_char}"
            f"{form}"
            f"{self.sep_char}"
            f"{src_right}"
        )

    def lemmatize(self, form, left_context=None, right_context=None):
        src_string = self.make_source(
            form=form,
            left_context=left_context,
            right_context=right_context,
        )

        src_ids = (
            [self.vocab.char2idx["<sos>"]]
            + self.vocab.encode(src_string)
            + [self.vocab.char2idx["<eos>"]]
        )

        src = torch.tensor(
            [src_ids],
            dtype=torch.long,
            device=self.device,
        )

        src_lens = torch.tensor(
            [len(src_ids)],
            dtype=torch.long,
            device=self.device,
        )

        return self.model.generate(
            src,
            src_lens,
            self.vocab,
            max_len=self.max_gen_len,
        )[0]

    def lemmatize_sentence(self, tokens):
        lemmas = []

        for index, token in enumerate(tokens):
            left_context = tokens[
                max(0, index - self.k_context):index
            ]

            right_context = tokens[
                index + 1:index + 1 + self.k_context
            ]

            lemma = self.lemmatize(
                form=token,
                left_context=left_context,
                right_context=right_context,
            )

            lemmas.append(lemma)

        return lemmas


def load_lemmatizer(model_id="ocscomdemo", device="cuda"):
    config_path, vocab_path, weights_path = download_model_files(
        model_id
    )

    with open(config_path, encoding="utf8") as file:
        config = json.load(file)

    with open(vocab_path, encoding="utf8") as file:
        vocab_data = json.load(file)

    vocab = Vocab(
        char2idx=vocab_data["char2idx"],
        idx2char=vocab_data["idx2char"],
    )

    expected_vocab_size = int(
        config.get("vocab_size", len(vocab.char2idx))
    )

    if expected_vocab_size != len(vocab.char2idx):
        raise ValueError(
            f"Vocabulary mismatch: config has "
            f"{expected_vocab_size}, vocab has "
            f"{len(vocab.char2idx)}"
        )

    model = LemmaModel(
        vocab_size=len(vocab.char2idx),
        char_emb_dim=int(config["char_emb_dim"]),
        hidden_size=int(config["hidden_size"]),
        drop_prob=float(config["drop_prob"]),
        num_heads=int(config["num_heads"]),
        max_gen_len=int(config.get("max_gen_len", 30)),
    )

    state = torch.load(
        weights_path,
        map_location="cpu",
        weights_only=True,
    )

    if isinstance(state, dict) and "state_dict" in state:
        state = state["state_dict"]

    model.load_state_dict(state, strict=True)
    model.to(device)
    model.eval()

    return OldSlavicLemmatizer(
        model=model,
        vocab=vocab,
        config=config,
        device=device,
    )