"""Prompt-level helpers for Tier-1 experiments on mlx_lm models: last-token hidden states per layer,
greedy generation with an intervention applied during the prompt pass only, and next-token NLL."""
from __future__ import annotations

import numpy as np

import mlx.core as mx
from mlx_lm.models.cache import make_prompt_cache

from wp_ai_0001.models.mlx_backend import MLXCapture, instrument


class Runner:
    def __init__(self, model, tokenizer, name: str) -> None:
        self.model, self.tok, self.name = model, tokenizer, name
        self.hooks = instrument(model)
        self.n_layers = len(model.model.layers)
        self.d = model.args.hidden_size

    def encode(self, text: str) -> list[int]:
        return self.tok.encode(text)

    def last_hidden(self, text: str, layers: list[int]) -> dict[int, np.ndarray]:
        cap = MLXCapture("last")
        names = [f"resid_post.{l}" for l in layers]
        for n in names:
            self.hooks.add(n, cap)
        try:
            self.model(mx.array([self.encode(text)]))
            mx.eval(*[cap.store[n] for n in names])
        finally:
            for n in names:
                self.hooks.remove(n, cap)
        return {l: np.array(cap.store[f"resid_post.{l}"][0].astype(mx.float32)) for l in layers}

    def generate(self, text: str, max_tokens: int = 6, intervention=None, layer: int | None = None) -> str:
        """Greedy decoding; ``intervention`` (callable(x, ctx)) is active only during the prompt pass at
        ``resid_post.{layer}``; subsequent steps run with hooks disabled (within-sequence persistence via KV)."""
        ids = self.encode(text)
        cache = make_prompt_cache(self.model)
        if intervention is not None:
            self.hooks.add(f"resid_post.{layer}", intervention)
        try:
            logits = self.model(mx.array([ids]), cache=cache)
        finally:
            if intervention is not None:
                self.hooks.remove(f"resid_post.{layer}", intervention)
        out = []
        prev = self.hooks.enabled
        self.hooks.enabled = False
        try:
            for _ in range(max_tokens):
                tok = int(mx.argmax(logits[0, -1, :]).item())
                out.append(tok)
                piece = self.tok.decode(out)
                if "\n" in piece:
                    break
                logits = self.model(mx.array([[tok]]), cache=cache)
        finally:
            self.hooks.enabled = prev
        return self.tok.decode(out)

    def next_token_nll(self, text: str, intervention=None, layer: int | None = None) -> float:
        """NLL of the last token of ``text`` given the preceding tokens, with an optional intervention at the
        prompt's last position (i.e. the position predicting that token)."""
        ids = self.encode(text)
        ctx, target = ids[:-1], ids[-1]
        if intervention is not None:
            self.hooks.add(f"resid_post.{layer}", intervention)
        try:
            logits = self.model(mx.array([ctx]))[0, -1, :].astype(mx.float32)
        finally:
            if intervention is not None:
                self.hooks.remove(f"resid_post.{layer}", intervention)
        lp = logits - mx.logsumexp(logits)
        return float(-lp[target].item())