"""Model interface contract + build dispatch + shared greedy decode."""
import torch
import torch.nn as nn

from src.config import Config
from src.data import pad_inputs


class PrimeModel(nn.Module):
    """Contract: forward(x, y_in) -> {"logits": (B,T_out,vocab), "halt_steps": (B,) or None}."""

    def forward(self, x, y_in):
        raise NotImplementedError

    def param_count(self) -> int:
        return sum(p.numel() for p in self.parameters())


def build_model(cfg: Config) -> PrimeModel:
    if cfg.model == "rnn":
        from src.models.rnn import TiedRNN
        return TiedRNN(cfg)
    if cfg.model == "transformer":
        from src.models.transformer import TransformerBaseline
        return TransformerBaseline(cfg)
    raise ValueError(f"unknown model: {cfg.model}")


@torch.no_grad()
def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int | None = None) -> torch.Tensor:
    """Autoregressive greedy decode of output digits. Returns (B, max_len) tokens (BOS stripped).

    Inputs are LEFT-padded to the global layout so positions match training exactly
    (batch/singleton invariance — codex BLOCKER fix)."""
    try:
        device = next(model.parameters()).device
    except StopIteration:
        device = x.device
    if x.device != device:
        x = x.to(device)
    x = pad_inputs(x, cfg)
    from src.data import _global_lengths
    _, out_max = _global_lengths(cfg)
    if max_len is None:
        max_len = max(getattr(cfg, "max_out_len", 6), out_max)
        if cfg.scratch_mode == "structured":
            # Probe-window traces ([range_end+1, +1000]) exceed the in-range layout
            # bound (measured: 633-token trace at n=1327 vs out_max=369). Cap at the
            # model's positional table instead so OOD traces decode fully.
            pos = getattr(model, "pos", None)
            if pos is not None:
                max_len = max(max_len, pos.num_embeddings - x.shape[1] - 1)
    B = x.shape[0]
    y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device)
    for _ in range(max_len):
        out = model(x, y_in)
        nxt = out["logits"][:, -1].argmax(-1)
        y_in = torch.cat([y_in, nxt[:, None]], dim=1)
    return y_in[:, 1:]