"""MinSparkForCausalLM: thin Transformers wrapper around the vendored Meiosis.

Exact semantics: identical to the bundled generate.py's generation
loop (EOS prefix once, truncate to last max_seq_len, effort -> loop count).
No KV cache (min-spark 1.1); right-padding is scoring-only; generation is
single-sequence (enforced in prepare_inputs_for_generation).
"""
from __future__ import annotations

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

try:
    from .configuration_minspark import MinSparkConfig  # remote-code: sibling in cache package
except ImportError:
    from configuration_minspark import MinSparkConfig  # direct import with staging on sys.path
try:
    from .meiosis import Meiosis, build_rope_cache  # remote-code: vendored sibling
except ImportError:
    from meiosis import Meiosis, build_rope_cache  # direct import with staging on sys.path

EFFORT_MAP = {"low": 2, "medium": 3, "high": 4}


class MinSparkForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = MinSparkConfig
    base_model_prefix = "model"
    main_input_name = "input_ids"
    supports_gradient_checkpointing = False
    _no_split_modules: list[str] = []

    def __init__(self, config: MinSparkConfig):
        super().__init__(config)
        self.model = Meiosis(config.to_meiosis())
        self.post_init()  # ties weights (no-op: output == input embedding)

    def get_input_embeddings(self) -> nn.Embedding:
        return self.model.embed

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.model.embed = value

    def get_output_embeddings(self) -> nn.Embedding:
        return self.model.embed  # tied: unembed reads embed.weight

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        effort: str | None = None,
        loops: int | None = None,
        past_key_values=None,
        use_cache: bool | None = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ) -> CausalLMOutputWithPast:
        if past_key_values is not None or use_cache:
            raise NotImplementedError(
                "KV cache is not implemented in min-spark; it arrives in 1.1. "
                "Set use_cache=False (the default)."
            )
        if input_ids.ndim != 2:
            raise ValueError(f"input_ids must be (B, T), got shape {tuple(input_ids.shape)}")
        if input_ids.shape[1] > self.config.max_seq_len:
            raise ValueError(
                f"seq_len {input_ids.shape[1]} > max {self.config.max_seq_len}; "
                "truncate the context or use generate (which truncates)."
            )
        self._validate_attention_mask(attention_mask, input_ids)
        loop_count = self._resolve_loops(effort, loops)
        self._ensure_buffers()

        logits = self.model(input_ids, loops=loop_count)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
            hidden_states=None,
            attentions=None,
        )

    def _resolve_loops(self, effort: str | None, loops: int | None) -> int:
        if loops is not None:
            if not isinstance(loops, int) or not (1 <= loops <= self.config.max_loops):
                raise ValueError(
                    f"loops must be an int in [1, {self.config.max_loops}], got {loops!r}"
                )
            return loops
        if effort is not None:
            if effort not in EFFORT_MAP:
                raise ValueError(f"effort must be one of {sorted(EFFORT_MAP)}, got {effort!r}")
            return EFFORT_MAP[effort]
        return EFFORT_MAP[self.config.effort]

    def _ensure_buffers(self) -> None:
        """Rebuild Meiosis's non-persistent buffers on first forward.

        from_pretrained constructs the model on torch.device('meta'), so
        build_rope_cache runs on meta tensors and yields garbage; transformers
        then restores only the persistent weights, never these non-persistent
        buffers. The garbage is not reliably non-finite (meta memory can be
        finite-but-wrong, e.g. 1e-21), so check the actual first-row value
        rather than finiteness, and rebuild unconditionally on first forward.
        Idempotent: runs once per instance."""
        if getattr(self, "_buffers_ok", False):
            return
        m = self.model
        cos, sin = build_rope_cache(m.config, m.config.max_seq_len)
        m.rope_cos.copy_(cos)
        m.rope_sin.copy_(sin)
        m.last_loop_rms.zero_()
        self._buffers_ok = True

    def _validate_attention_mask(
        self, attention_mask: torch.Tensor | None, input_ids: torch.Tensor
    ) -> None:
        if attention_mask is None:
            return
        if tuple(attention_mask.shape) != tuple(input_ids.shape):
            raise ValueError(
                f"attention_mask shape {tuple(attention_mask.shape)} != "
                f"input_ids shape {tuple(input_ids.shape)}"
            )
        mask = attention_mask.bool()
        # leading zeros = left padding (any row whose FIRST position is masked out)
        if mask.shape[1] >= 1 and (~mask[:, 0]).any():
            raise ValueError(
                "left-padded batches are not supported; pad to the right or run single-sequence"
            )
        # interior gap: a 0 followed later by a 1
        if mask.shape[1] >= 2 and (mask[:, 1:].long() - mask[:, :-1].long() > 0).any():
            raise ValueError(
                "attention_mask must be ones or a contiguous ones-then-zeros suffix; "
                "interior gaps are not supported"
            )

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        effort: str | None = None,
        loops: int | None = None,
        **kwargs,
    ):
        """Build the next forward's inputs. Returns exactly these four keys so
        generation machinery (cache_position, position_ids, use_cache) is never
        echoed into forward, which has no **kwargs. EOS is prepended BEFORE
        truncation (generate.py parity — it drops off prompts >max_seq_len);
        a supplied attention_mask is extended/truncated in lockstep so its length
        always matches the returned input_ids (forward validates mask shape)."""
        if input_ids.shape[0] != 1:
            raise ValueError(
                "batched generation is not supported; run single-sequence generation "
                "or right-padded scoring through forward"
            )
        ids = input_ids
        mask = attention_mask
        if self.config.doc_mask_eos is not None:
            eos = self.config.doc_mask_eos
            ids = torch.cat(
                [torch.full((1, 1), eos, dtype=ids.dtype, device=ids.device), ids], dim=1
            )
            if mask is not None:
                # The prepended EOS is a real position: keep the mask in sync.
                mask = torch.cat(
                    [torch.ones((1, 1), dtype=mask.dtype, device=mask.device), mask], dim=1
                )
        if ids.shape[1] > self.config.max_seq_len:
            ids = ids[:, -self.config.max_seq_len:]
            if mask is not None:
                mask = mask[:, -self.config.max_seq_len:]
        return {
            "input_ids": ids,
            "attention_mask": mask,
            "effort": effort,
            "loops": loops,
        }

    def _reorder_cache(self, past_key_values, beam_idx):
        return past_key_values  # no cache