import json
from functools import lru_cache
from pathlib import Path

import torch
from huggingface_hub import hf_hub_download

try:
    from model import LemmaModel, Vocab
except ImportError:
    from .model import LemmaModel, Vocab


MODEL_REPO_ID = "usmannawaz/oldslaviclemma"
MODEL_ROOT = "oldslaviclemma212"


@lru_cache(maxsize=1)
def load_registry(registry_path="models_registry.json"):
    
    registry_path = Path(registry_path)
    with registry_path.open(encoding="utf8") as f:
        return json.load(f)


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}{self.sep_char}{form}{self.sep_char}{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 i, token in enumerate(tokens):
            left_context = tokens[max(0, i - self.k_context):i]
            right_context = tokens[i + 1:i + 1 + self.k_context]

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

            lemmas.append(lemma)

        return lemmas


@lru_cache(maxsize=3)
def load_lemmatizer(model_id, device=None):
    
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    device = torch.device(device)

    registry = load_registry("models_registry.json")
    if model_id not in registry:
        raise KeyError(f"Model id not found in registry: {model_id}")

    item = registry[model_id]
    folder = item["folder"]

    config_path = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        repo_type="model",
        filename=f"{MODEL_ROOT}/{folder}/{item['config_file']}",
    )

    vocab_path = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        repo_type="model",
        filename=f"{MODEL_ROOT}/{folder}/{item['vocab_file']}",
    )

    weights_path = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        repo_type="model",
        filename=f"{MODEL_ROOT}/{folder}/{item['model_file']}",
    )

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

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

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

    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)),
    ).to(device)

    state = torch.load(weights_path, map_location=device)
    model.load_state_dict(state)
    model.eval()

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