import dataclasses

import torch
import torch.nn.functional as F
from transformers import GenerationMixin, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_pragya import PragyaConfig
from .model import GPT, GPTConfig
# The Hub loader only copies sibling files named in "from .name import" form. model.py
# falls back to "from . import config/kernel/ngram", so import them this way too.
from .config import *  # noqa: F401,F403
from .kernel import *  # noqa: F401,F403
from .ngram import *  # noqa: F401,F403

# Config field names that differ between PragyaConfig and GPTConfig.
_RENAMED = {"top_k": "gpt_top_k"}
# Fields GPTConfig stores as tuples, which PragyaConfig stores as lists.
_TUPLE_FIELDS = ("engram_layer_ids", "engram_vocab_size")


def to_gpt_config(cfg: PragyaConfig) -> GPTConfig:
    """Build a GPTConfig from every field, so no value comes from config.py defaults."""
    kwargs = {}
    for f in dataclasses.fields(GPTConfig):
        value = getattr(cfg, _RENAMED.get(f.name, f.name))
        if f.name in _TUPLE_FIELDS:
            value = tuple(value)
        kwargs[f.name] = value
    return GPTConfig(**kwargs)


def _non_persistent_buffers(module):
    """{name: tensor} for every non-persistent buffer in a torch module tree."""
    out = {}
    for prefix, sub in module.named_modules():
        for buf_name in sub._non_persistent_buffers_set:
            full = f"{prefix}.{buf_name}" if prefix else buf_name
            out[full] = getattr(sub, buf_name)
    return out


class PragyaForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = PragyaConfig
    base_model_prefix = "model"
    main_input_name = "input_ids"
    supports_gradient_checkpointing = False
    # wte.weight and lm_head.weight are one tensor; save only lm_head.weight.
    _tied_weights_keys = {"model.transformer.wte.weight": "model.lm_head.weight"}

    def __init__(self, config: PragyaConfig):
        super().__init__(config)
        if config.use_engram and config.engram_compressed_vocab <= 0:
            raise ValueError(
                "engram_compressed_vocab must be set before loading. It is computed from the "
                "tokenizer when the model is first built, so the export step must record it."
            )
        self.model = GPT(to_gpt_config(config))
        self.post_init()

    def _move_missing_keys_from_meta_to_device(self, missing_keys, device_map, device_mesh, hf_quantizer):
        # Transformers replaces every non-persistent buffer (RoPE tables, XSA masks, Engram
        # `mods`) with uninitialised memory when loading. They are derived from the config,
        # so rebuild them from a freshly constructed GPT.
        super()._move_missing_keys_from_meta_to_device(missing_keys, device_map, device_mesh, hf_quantizer)
        fresh = GPT(to_gpt_config(self.config))
        own = _non_persistent_buffers(self.model)
        with torch.no_grad():
            for name, buf in _non_persistent_buffers(fresh).items():
                own[name].copy_(buf)
        del fresh
        # wte and lm_head are one tensor. Checkpoints store one of the two, so tie from the
        # one that was actually loaded, not from whichever happens to be first.
        if "model.lm_head.weight" in missing_keys:
            self.model.lm_head.weight = self.model.transformer.wte.weight
        else:
            self.model.transformer.wte.weight = self.model.lm_head.weight

    def _init_weights(self, module):
        # Weights come from GPT.__init__ (Engram tables, zero-init attention gate, tied
        # embeddings) or from the checkpoint. Transformers' default would overwrite them.
        pass

    def get_input_embeddings(self):
        return self.model.transformer.wte

    def get_output_embeddings(self):
        return self.model.lm_head

    def forward(self, input_ids=None, attention_mask=None, labels=None,
                logits_to_keep=0, **kwargs):
        if attention_mask is not None and not bool(torch.all(attention_mask == 1)):
            # The model has no attention-mask support. Padded batches would silently
            # attend to pad tokens, so refuse rather than return wrong results.
            raise NotImplementedError(
                "attention_mask with padding is not supported. Run one sequence at a time "
                "or pass unpadded batches of equal length."
            )
        if labels is not None:
            # Training path: GPT computes the loss itself.
            logits, loss = self.model(input_ids, targets=labels)
            return CausalLMOutputWithPast(loss=loss, logits=logits)

        # Inference path. logits_to_keep=0 means logits for every position (B, T, V),
        # which lm-evaluation-harness and loglikelihood scoring need. logits_to_keep=k
        # returns the last k positions. Hidden states are captured after ln_f, and the
        # projection is the same F.linear(x, lm_head.weight) GPT uses internally.
        captured = {}
        handle = self.model.transformer.ln_f.register_forward_hook(
            lambda module, inp, out: captured.__setitem__("x", out))
        try:
            self.model(input_ids)
        finally:
            handle.remove()
        x = captured["x"]
        if logits_to_keep:
            x = x[:, -logits_to_keep:, :]
        logits = F.linear(x, self.model.lm_head.weight)
        return CausalLMOutputWithPast(logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        # No KV cache: GPT keeps its own cache format, so every step runs on the full sequence.
        return {"input_ids": input_ids}