"""
Generation and prediction tools: generate_text, predict_next_token,
tokenize, logit_lens, track_token, embedding_neighbors.

These tools expose the model's basic I/O capabilities that are essential
for an AI agent to understand *what* a model produces, complementing
the interpretability tools that explain *why* it produces it.
"""

import asyncio
import logging
from typing import Any

import mlx.core as mx
from pydantic import BaseModel, Field

from .._serialize import to_pylist
from ..errors import ToolError, make_error
from ..model_state import ModelState
from ..server import mcp

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Result models
# ---------------------------------------------------------------------------


class GenerateResult(BaseModel):
    """Result from generate_text."""

    prompt: str
    output: str
    num_tokens: int
    temperature: float
    max_new_tokens: int


class PredictionEntry(BaseModel):
    """A single next-token prediction."""

    token: str
    token_id: int
    probability: float
    log_probability: float


class PredictResult(BaseModel):
    """Result from predict_next_token."""

    prompt: str
    num_input_tokens: int
    predictions: list[PredictionEntry]


class TokenInfo(BaseModel):
    """Info for a single token."""

    position: int
    token_id: int
    token_text: str


class TokenizeResult(BaseModel):
    """Result from tokenize."""

    text: str
    num_tokens: int
    tokens: list[TokenInfo]
    token_ids: list[int]


class LayerPredictionEntry(BaseModel):
    """Top-k predictions at a single layer."""

    layer: int
    top_tokens: list[str]
    top_probabilities: list[float]
    top_token_ids: list[int]


class LogitLensResult(BaseModel):
    """Result from logit_lens."""

    prompt: str
    token_position: int
    token_text: str
    num_layers_analyzed: int
    predictions: list[LayerPredictionEntry]
    summary: dict = Field(..., description="Where the final prediction first emerges.")


class TokenLayerEntry(BaseModel):
    """Probability and rank of a tracked token at one layer."""

    layer: int
    probability: float
    rank: int
    is_top1: bool


class TrackTokenResult(BaseModel):
    """Result from track_token."""

    prompt: str
    target_token: str
    target_token_id: int
    token_position: int
    token_text: str
    layers: list[TokenLayerEntry]
    emergence_layer: int | None = Field(
        None, description="First layer where target is top-1 prediction."
    )
    peak_layer: int
    peak_probability: float


class CandidateLayerEntry(BaseModel):
    """Per-layer probability and rank for one candidate token."""

    layer: int
    probability: float
    rank: int
    is_top1: bool


class CandidateTrajectory(BaseModel):
    """A single candidate token's trajectory across layers."""

    token: str
    token_id: int
    layers: list[CandidateLayerEntry]
    emergence_layer: int | None = Field(
        None, description="First layer where this candidate is top-1."
    )
    peak_layer: int
    peak_probability: float


class CrossingEvent(BaseModel):
    """A layer where the lead candidate changes."""

    layer: int
    previous_leader: str
    new_leader: str
    new_leader_probability: float


class TrackRaceResult(BaseModel):
    """Result from track_race."""

    prompt: str
    token_position: int
    token_text: str
    num_candidates: int
    candidates: list[CandidateTrajectory]
    crossings: list[CrossingEvent]
    final_winner: str
    final_probability: float


class NeighborToken(BaseModel):
    """A token similar to the query in embedding space."""

    token: str
    token_id: int
    cosine_similarity: float


class EmbeddingNeighborsResult(BaseModel):
    """Result from embedding_neighbors."""

    query_token: str
    query_token_id: int
    resolved_form: str = Field(
        ..., description="Actual form used (may differ from input due to space-prefix)."
    )
    embedding_dim: int
    vocab_size: int
    top_k: int
    neighbors: list[NeighborToken]
    self_similarity: float = Field(
        ..., description="Cosine similarity of the token with itself (sanity check, ~1.0)."
    )


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def generate_text(
    prompt: str,
    max_new_tokens: int = 100,
    temperature: float = 0.0,
) -> dict:
    """
    Generate text from the loaded model.

    Uses greedy decoding when temperature is 0, otherwise samples.
    This is the basic "what does the model say?" tool -- essential for
    seeing actual model outputs before and after interpretability work.

    Args:
        prompt:         Input text.
        max_new_tokens: Maximum tokens to generate (default: 100).
        temperature:    Sampling temperature. 0 = greedy (default).
    """
    state = ModelState.get()
    if not state.is_loaded:
        return make_error(
            ToolError.MODEL_NOT_LOADED,
            "Call load_model() first.",
            "generate_text",
        )

    if max_new_tokens < 1 or max_new_tokens > 1000:
        return make_error(
            ToolError.INVALID_INPUT,
            f"max_new_tokens must be 1-1000, got {max_new_tokens}.",
            "generate_text",
        )

    try:
        from .._generate import generate_text as _gen

        output, num_tokens = await asyncio.to_thread(
            _gen,
            state.model,
            state.tokenizer,
            prompt,
            max_new_tokens,
            temperature,
        )

        result = GenerateResult(
            prompt=prompt,
            output=output,
            num_tokens=num_tokens,
            temperature=temperature,
            max_new_tokens=max_new_tokens,
        )
        return result.model_dump()

    except Exception as e:
        logger.exception("generate_text failed")
        return make_error(ToolError.GENERATION_FAILED, str(e), "generate_text")


@mcp.tool(read_only_hint=True, idempotent_hint=True)
async def predict_next_token(
    prompt: str,
    top_k: int = 10,
) -> dict:
    """
    Get the model's top-k next-token predictions with probabilities.

    Shows what the model thinks comes next without generating text.
    Useful for understanding model confidence and comparing how
    different prompts affect the prediction distribution.

    Args:
        prompt: Input text.
        top_k:  Number of top predictions to return (default: 10).
    """
    state = ModelState.get()
    if not state.is_loaded:
        return make_error(
            ToolError.MODEL_NOT_LOADED,
            "Call load_model() first.",
            "predict_next_token",
        )

    if top_k < 1 or top_k > 100:
        return make_error(
            ToolError.INVALID_INPUT,
            f"top_k must be 1-100, got {top_k}.",
            "predict_next_token",
        )

    try:
        predictions = await asyncio.to_thread(
            _predict_next,
            state.model,
            state.tokenizer,
            prompt,
            top_k,
        )
        return predictions

    except Exception as e:
        logger.exception("predict_next_token failed")
        return make_error(ToolError.GENERATION_FAILED, str(e), "predict_next_token")


def _predict_next(model: Any, tokenizer: Any, prompt: str, top_k: int) -> dict:
    """Compute top-k next-token predictions."""
    input_ids = mx.array(tokenizer.encode(prompt, add_special_tokens=True))
    num_input = input_ids.shape[-1]

    logits = model(input_ids[None, :] if input_ids.ndim == 1 else input_ids)
    if isinstance(logits, tuple):
        logits = logits[0]
    if hasattr(logits, "logits"):
        logits = logits.logits

    next_logits = logits[0, -1, :] if logits.ndim == 3 else logits[-1, :]

    probs = mx.softmax(next_logits, axis=-1)
    log_probs = mx.log(probs + 1e-10)

    sorted_indices = mx.argsort(probs)[::-1][:top_k]
    mx.eval(sorted_indices, probs, log_probs)

    top_ids = to_pylist(sorted_indices)
    top_probs = to_pylist(probs[sorted_indices])
    top_log_probs = to_pylist(log_probs[sorted_indices])

    entries = []
    for tid, p, lp in zip(top_ids, top_probs, top_log_probs):
        token_text = tokenizer.decode([tid])
        entries.append(
            PredictionEntry(
                token=token_text,
                token_id=tid,
                probability=round(p, 6),
                log_probability=round(lp, 4),
            )
        )

    result = PredictResult(
        prompt=prompt,
        num_input_tokens=int(num_input),
        predictions=entries,
    )
    return result.model_dump()


@mcp.tool(read_only_hint=True, idempotent_hint=True)
async def tokenize(
    text: str,
) -> dict:
    """
    Show how text is tokenized by the loaded model's tokenizer.

    Essential for understanding attention patterns (which token is being
    attended to) and debugging multi-token words. Returns token IDs,
    decoded text for each token, and positions.

    Args:
        text: Input text to tokenize.
    """
    state = ModelState.get()
    if not state.is_loaded:
        return make_error(
            ToolError.MODEL_NOT_LOADED,
            "Call load_model() first.",
            "tokenize",
        )

    try:
        token_ids = state.tokenizer.encode(text, add_special_tokens=True)

        tokens = []
        for i, tid in enumerate(token_ids):
            decoded = state.tokenizer.decode([tid])
            tokens.append(
                TokenInfo(
                    position=i,
                    token_id=tid,
                    token_text=decoded,
                )
            )

        result = TokenizeResult(
            text=text,
            num_tokens=len(token_ids),
            tokens=tokens,
            token_ids=token_ids,
        )
        return result.model_dump()

    except Exception as e:
        logger.exception("tokenize failed")
        return make_error(ToolError.GENERATION_FAILED, str(e), "tokenize")


@mcp.tool(read_only_hint=True, idempotent_hint=True)
async def logit_lens(
    prompt: str,
    layers: list[int] | None = None,
    top_k: int = 5,
    token_position: int = -1,
) -> dict:
    """
    Apply the "logit lens" technique: project hidden states from each
    layer back to vocabulary space to see how the model's prediction
    evolves through its depth.

    Shows what the model would predict if computation stopped at each
    layer. Critical for understanding when specific knowledge emerges
    (e.g., "at which layer does the model start predicting the correct
    translation token?").

    Args:
        prompt:         Input text.
        layers:         Layer indices to analyze. None = sample ~10 layers.
        top_k:          Number of top predictions per layer (default: 5).
        token_position: Which token to analyze (-1 = last, default).
    """
    state = ModelState.get()
    if not state.is_loaded:
        return make_error(
            ToolError.MODEL_NOT_LOADED,
            "Call load_model() first.",
            "logit_lens",
        )

    if top_k < 1 or top_k > 100:
        return make_error(
            ToolError.INVALID_INPUT,
            f"top_k must be 1-100, got {top_k}.",
            "logit_lens",
        )

    num_layers = state.metadata.num_layers

    if layers is None:
        # Sample ~10 layers across the model
        if num_layers <= 10:
            layers = list(range(num_layers))
        else:
            step = max(1, num_layers // 10)
            layers = list(range(0, num_layers, step))
            if (num_layers - 1) not in layers:
                layers.append(num_layers - 1)

    out_of_range = [lay for lay in layers if lay < 0 or lay >= num_layers]
    if out_of_range:
        return make_error(
            ToolError.LAYER_OUT_OF_RANGE,
            f"Layers {out_of_range} out of range [0, {num_layers - 1}].",
            "logit_lens",
        )

    try:
        result = await asyncio.to_thread(
            _logit_lens_impl,
            state.model,
            state.config,
            state.tokenizer,
            prompt,
            sorted(layers),
            top_k,
            token_position,
        )
        return result

    except Exception as e:
        logger.exception("logit_lens failed")
        return make_error(ToolError.EXTRACTION_FAILED, str(e), "logit_lens")


def _logit_lens_impl(
    model: Any,
    config: Any,
    tokenizer: Any,
    prompt: str,
    layers: list[int],
    top_k: int,
    token_position: int,
) -> dict:
    """Run logit lens analysis with calibrated norm+head projection.

    Uses _get_lm_projection (handles tied-embedding models like Gemma 3)
    and _norm_project (applies final_norm before projection) instead of
    hooks.get_layer_logits which can use the wrong lm_head.
    """
    from chuk_lazarus.introspection.hooks import CaptureConfig, ModelHooks

    from .residual_tools import _get_lm_projection, _norm_project

    input_ids = mx.array(tokenizer.encode(prompt, add_special_tokens=True))
    num_tokens = input_ids.shape[-1]

    # Resolve token text at the position
    ids_list = to_pylist(input_ids)
    pos = token_position if token_position >= 0 else num_tokens + token_position
    pos = max(0, min(pos, num_tokens - 1))
    token_text = tokenizer.decode([ids_list[pos]])

    hooks = ModelHooks(model, model_config=config)
    hooks.configure(
        CaptureConfig(
            layers=layers,
            capture_hidden_states=True,
        )
    )
    hooks.forward(input_ids)
    mx.eval(hooks.state.hidden_states)

    # Calibrated projection: handles tied-embedding models correctly
    lm_head = _get_lm_projection(model)
    helper = ModelHooks(model, model_config=config)
    final_norm = helper._get_final_norm()

    predictions = []
    for layer_idx in layers:
        h = hooks.state.hidden_states.get(layer_idx)
        if h is None:
            continue

        # Extract hidden state at the target position
        if h.ndim == 3:
            vec = h[0, token_position, :]
        elif h.ndim == 2:
            vec = h[token_position, :]
        else:
            vec = h

        # Project through final_norm + lm_head
        pos_logits = _norm_project(final_norm, lm_head, vec)

        probs = mx.softmax(pos_logits, axis=-1)
        sorted_indices = mx.argsort(probs)[::-1][:top_k]
        mx.eval(sorted_indices, probs)

        top_ids = to_pylist(sorted_indices)
        top_probs = to_pylist(probs[sorted_indices])
        top_tokens = [tokenizer.decode([tid]) for tid in top_ids]

        predictions.append(
            LayerPredictionEntry(
                layer=layer_idx,
                top_tokens=top_tokens,
                top_probabilities=[round(p, 6) for p in top_probs],
                top_token_ids=top_ids,
            )
        )

    # Summary: when does the final prediction first emerge?
    final_token = predictions[-1].top_tokens[0] if predictions else None
    emergence_layer = None
    if final_token is not None:
        for pred in predictions:
            if pred.top_tokens[0] == final_token:
                emergence_layer = pred.layer
                break

    summary = {
        "final_prediction": final_token,
        "emergence_layer": emergence_layer,
        "total_layers": len(predictions),
    }

    result = LogitLensResult(
        prompt=prompt,
        token_position=token_position,
        token_text=token_text,
        num_layers_analyzed=len(predictions),
        predictions=predictions,
        summary=summary,
    )
    return result.model_dump()


@mcp.tool(read_only_hint=True, idempotent_hint=True)
async def track_token(
    prompt: str,
    token: str,
    layers: list[int] | None = None,
    token_position: int = -1,
) -> dict:
    """
    Track a specific token's probability and rank across layers.

    Uses the logit lens technique to project hidden states at each layer
    back to vocabulary space, then finds the target token's probability
    and rank. Answers "at which layer does the model start predicting
    this specific token?"

    Args:
        prompt:         Input text.
        token:          Target token to track (string, e.g. "Paris").
        layers:         Layer indices to analyze. None = sample ~10 layers.
        token_position: Which input token to analyze (-1 = last, default).
    """
    state = ModelState.get()
    if not state.is_loaded:
        return make_error(
            ToolError.MODEL_NOT_LOADED,
            "Call load_model() first.",
            "track_token",
        )

    num_layers = state.metadata.num_layers

    if layers is None:
        if num_layers <= 10:
            layers = list(range(num_layers))
        else:
            step = max(1, num_layers // 10)
            layers = list(range(0, num_layers, step))
            if (num_layers - 1) not in layers:
                layers.append(num_layers - 1)

    out_of_range = [lay for lay in layers if lay < 0 or lay >= num_layers]
    if out_of_range:
        return make_error(
            ToolError.LAYER_OUT_OF_RANGE,
            f"Layers {out_of_range} out of range [0, {num_layers - 1}].",
            "track_token",
        )

    # Resolve target token ID
    token_ids = state.tokenizer.encode(token, add_special_tokens=False)
    if not token_ids:
        return make_error(
            ToolError.INVALID_INPUT,
            f"Could not encode token {token!r}.",
            "track_token",
        )
    target_token_id = token_ids[0]

    try:
        result = await asyncio.to_thread(
            _track_token_impl,
            state.model,
            state.config,
            state.tokenizer,
            prompt,
            sorted(layers),
            target_token_id,
            token,
            token_position,
        )
        return result

    except Exception as e:
        logger.exception("track_token failed")
        return make_error(ToolError.EXTRACTION_FAILED, str(e), "track_token")


def _track_token_impl(
    model: Any,
    config: Any,
    tokenizer: Any,
    prompt: str,
    layers: list[int],
    target_token_id: int,
    target_token: str,
    token_position: int,
) -> dict:
    """Track a specific token's probability across layers.

    Uses calibrated norm+head projection (handles tied-embedding models).
    """
    from chuk_lazarus.introspection.hooks import CaptureConfig, ModelHooks

    from .residual_tools import _get_lm_projection, _norm_project

    input_ids = mx.array(tokenizer.encode(prompt, add_special_tokens=True))
    num_tokens = input_ids.shape[-1]

    # Resolve token text at position
    ids_list = to_pylist(input_ids)
    pos = token_position if token_position >= 0 else num_tokens + token_position
    pos = max(0, min(pos, num_tokens - 1))
    token_text = tokenizer.decode([ids_list[pos]])

    hooks = ModelHooks(model, model_config=config)
    hooks.configure(
        CaptureConfig(
            layers=layers,
            capture_hidden_states=True,
        )
    )
    hooks.forward(input_ids)
    mx.eval(hooks.state.hidden_states)

    # Calibrated projection: handles tied-embedding models correctly
    lm_head = _get_lm_projection(model)
    helper = ModelHooks(model, model_config=config)
    final_norm = helper._get_final_norm()

    entries = []
    emergence_layer = None
    peak_layer = layers[0]
    peak_probability = 0.0

    for layer_idx in layers:
        h = hooks.state.hidden_states.get(layer_idx)
        if h is None:
            continue

        # Extract hidden state at the target position
        if h.ndim == 3:
            vec = h[0, token_position, :]
        elif h.ndim == 2:
            vec = h[token_position, :]
        else:
            vec = h

        # Project through final_norm + lm_head
        pos_logits = _norm_project(final_norm, lm_head, vec)

        probs = mx.softmax(pos_logits, axis=-1)

        # Get target token's probability
        target_prob = float(probs[target_token_id])

        # Get rank (how many tokens have higher probability)
        rank = int(mx.sum(probs > probs[target_token_id]))

        is_top1 = rank == 0
        if is_top1 and emergence_layer is None:
            emergence_layer = layer_idx

        if target_prob > peak_probability:
            peak_probability = target_prob
            peak_layer = layer_idx

        entries.append(
            TokenLayerEntry(
                layer=layer_idx,
                probability=round(target_prob, 6),
                rank=rank,
                is_top1=is_top1,
            )
        )

    result = TrackTokenResult(
        prompt=prompt,
        target_token=target_token,
        target_token_id=target_token_id,
        token_position=token_position,
        token_text=token_text,
        layers=entries,
        emergence_layer=emergence_layer,
        peak_layer=peak_layer,
        peak_probability=round(peak_probability, 6),
    )
    return result.model_dump()


@mcp.tool(read_only_hint=True, idempotent_hint=True)
async def embedding_neighbors(
    token: str,
    top_k: int = 20,
) -> dict:
    """Find the most similar tokens in embedding space.

    Uses cosine similarity between embedding vectors.  No forward pass
    needed -- operates directly on the static embedding matrix.

    Useful for understanding vocabulary structure: which tokens the
    model considers similar before any context is applied.

    Args:
        token:  Token string to find neighbours for (e.g. ``"Paris"``).
        top_k:  Number of nearest neighbours to return (default: 20).
    """
    state = ModelState.get()
    if not state.is_loaded:
        return make_error(
            ToolError.MODEL_NOT_LOADED,
            "No model loaded. Call load_model first.",
            "embedding_neighbors",
        )

    if top_k < 1:
        return make_error(
            ToolError.INVALID_INPUT,
            f"top_k must be >= 1, got {top_k}.",
            "embedding_neighbors",
        )
    top_k = min(top_k, 100)

    try:
        result = await asyncio.to_thread(
            _embedding_neighbors_impl,
            state.model,
            state.tokenizer,
            token,
            top_k,
        )
        return result

    except Exception as e:
        logger.exception("embedding_neighbors failed")
        return make_error(ToolError.EXTRACTION_FAILED, str(e), "embedding_neighbors")


def _embedding_neighbors_impl(
    model: Any,
    tokenizer: Any,
    token: str,
    top_k: int,
) -> dict:
    """Synchronous implementation of embedding_neighbors."""
    # Resolve token to a single token id
    # Try bare encoding first, then space-prefixed
    token_id = None
    resolved_form = token

    for variant in (token, " " + token):
        ids = tokenizer.encode(variant, add_special_tokens=False)
        if ids:
            token_id = ids[0]
            resolved_form = tokenizer.decode([token_id])
            break

    if token_id is None:
        return make_error(
            ToolError.INVALID_INPUT,
            f"Could not encode token {token!r}.",
            "embedding_neighbors",
        )

    # Get embedding matrix — handle chuk-lazarus wrappers where
    # embed_tokens.weight may be nn.Embedding rather than mx.array
    inner = getattr(model, "model", model)
    embed = getattr(inner, "embed_tokens", None)
    W = None
    if embed is not None:
        w = getattr(embed, "weight", None)
        if isinstance(w, mx.array):
            W = w
        elif w is not None:
            # nn.Embedding wrapper — actual array one level deeper
            ww = getattr(w, "weight", None)
            if isinstance(ww, mx.array):
                W = ww

    if W is None:
        return make_error(
            ToolError.EXTRACTION_FAILED,
            "Could not access embedding matrix.",
            "embedding_neighbors",
        )

    vocab_size, hidden_dim = W.shape

    # Target vector
    target_vec = W[token_id]  # [hidden_dim]

    # Compute all cosine similarities in one matrix-vector multiply
    # sims = (W @ target) / (||W|| * ||target||)
    dots = W @ target_vec  # [vocab_size]
    norms = mx.sqrt(mx.sum(W * W, axis=1))  # [vocab_size]
    target_norm = mx.sqrt(mx.sum(target_vec * target_vec))

    # Avoid division by zero
    denom = norms * target_norm + 1e-8
    sims = dots / denom

    # Clamp for bfloat16 safety
    sims = mx.clip(sims, -1.0, 1.0)
    mx.eval(sims)

    # Self similarity
    self_sim = float(sims[token_id].item())

    # Sort descending and take top_k+1 (to skip self)
    sorted_indices = mx.argsort(sims)[::-1]
    mx.eval(sorted_indices)
    sorted_list = to_pylist(sorted_indices)

    neighbors: list[NeighborToken] = []
    for idx in sorted_list:
        if idx == token_id:
            continue
        neighbors.append(
            NeighborToken(
                token=tokenizer.decode([idx]),
                token_id=idx,
                cosine_similarity=round(float(sims[idx].item()), 6),
            )
        )
        if len(neighbors) >= top_k:
            break

    result = EmbeddingNeighborsResult(
        query_token=token,
        query_token_id=token_id,
        resolved_form=resolved_form,
        embedding_dim=hidden_dim,
        vocab_size=vocab_size,
        top_k=top_k,
        neighbors=neighbors,
        self_similarity=round(self_sim, 6),
    )
    return result.model_dump()


# ---------------------------------------------------------------------------
# Tool: track_race
# ---------------------------------------------------------------------------


@mcp.tool(read_only_hint=True, idempotent_hint=True)
async def track_race(
    prompt: str,
    candidates: list[str],
    layers: list[int] | None = None,
    token_position: int = -1,
) -> dict:
    """
    Track multiple candidate tokens' probabilities across layers.

    Runs a single forward pass and projects hidden states at each layer
    to vocabulary space, then extracts probability and rank for each
    candidate token. Detects "crossings" where the leading candidate
    changes. Useful for comparing how competing predictions evolve
    through the model (e.g. "Canberra" vs "Sydney" vs "Melbourne").

    Args:
        prompt:         Input text.
        candidates:     List of candidate tokens to race (2–20).
        layers:         Layer indices to analyze. None = sample ~12 layers.
        token_position: Which input token to analyze (-1 = last, default).
    """
    state = ModelState.get()
    if not state.is_loaded:
        return make_error(
            ToolError.MODEL_NOT_LOADED,
            "Call load_model() first.",
            "track_race",
        )

    if len(candidates) < 2:
        return make_error(
            ToolError.INVALID_INPUT,
            "At least 2 candidate tokens required.",
            "track_race",
        )

    if len(candidates) > 20:
        return make_error(
            ToolError.INVALID_INPUT,
            "Maximum 20 candidates allowed.",
            "track_race",
        )

    num_layers = state.metadata.num_layers

    if layers is None:
        if num_layers <= 12:
            layers = list(range(num_layers))
        else:
            step = max(1, num_layers // 12)
            layers = list(range(0, num_layers, step))
            if (num_layers - 1) not in layers:
                layers.append(num_layers - 1)

    out_of_range = [lay for lay in layers if lay < 0 or lay >= num_layers]
    if out_of_range:
        return make_error(
            ToolError.LAYER_OUT_OF_RANGE,
            f"Layers {out_of_range} out of range [0, {num_layers - 1}].",
            "track_race",
        )

    try:
        result = await asyncio.to_thread(
            _track_race_impl,
            state.model,
            state.config,
            state.tokenizer,
            prompt,
            candidates,
            sorted(layers),
            token_position,
        )
        return result
    except ValueError as exc:
        return make_error(ToolError.INVALID_INPUT, str(exc), "track_race")
    except Exception as e:
        logger.exception("track_race failed")
        return make_error(ToolError.EXTRACTION_FAILED, str(e), "track_race")


def _track_race_impl(
    model: Any,
    config: Any,
    tokenizer: Any,
    prompt: str,
    candidates: list[str],
    layers: list[int],
    token_position: int,
) -> dict:
    """Track multiple tokens' probabilities across layers.

    Uses calibrated norm+head projection (handles tied-embedding models).
    """
    from chuk_lazarus.introspection.hooks import CaptureConfig, ModelHooks

    from .residual_tools import _get_lm_projection, _norm_project

    input_ids = mx.array(tokenizer.encode(prompt, add_special_tokens=True))
    num_tokens = input_ids.shape[-1]

    # Resolve token text at position
    ids_list = to_pylist(input_ids)
    pos = token_position if token_position >= 0 else num_tokens + token_position
    pos = max(0, min(pos, num_tokens - 1))
    token_text = tokenizer.decode([ids_list[pos]])

    # Forward pass — capture hidden states at all requested layers
    hooks = ModelHooks(model, model_config=config)
    hooks.configure(
        CaptureConfig(
            layers=layers,
            capture_hidden_states=True,
        )
    )
    hooks.forward(input_ids)
    mx.eval(hooks.state.hidden_states)

    # Calibrated projection
    lm_head = _get_lm_projection(model)
    helper = ModelHooks(model, model_config=config)
    final_norm = helper._get_final_norm()

    # Get logits from last layer for candidate resolution
    last_layer = layers[-1]
    h_last = hooks.state.hidden_states.get(last_layer)
    if h_last is None:
        raise ValueError(f"No hidden state captured at layer {last_layer}")
    if h_last.ndim == 3:
        vec_last = h_last[0, token_position, :]
    elif h_last.ndim == 2:
        vec_last = h_last[token_position, :]
    else:
        vec_last = h_last
    full_logits = _norm_project(final_norm, lm_head, vec_last)

    # Resolve all candidate tokens
    resolved: list[tuple[int, str]] = []
    for cand in candidates:
        cand_ids: list[int] = []
        for variant in (cand, " " + cand):
            ids = tokenizer.encode(variant, add_special_tokens=False)
            if ids:
                cand_ids.append(ids[0])
        # Deduplicate
        seen: set[int] = set()
        unique: list[int] = []
        for c in cand_ids:
            if c not in seen:
                seen.add(c)
                unique.append(c)
        if not unique:
            raise ValueError(f"Could not encode candidate token {cand!r}")
        tid = max(unique, key=lambda t: float(full_logits[t].item()))
        resolved.append((tid, tokenizer.decode([tid])))

    # Track across layers
    trajectories: list[dict] = []
    for tid, tok_text in resolved:
        trajectories.append(
            {
                "token": tok_text,
                "token_id": tid,
                "entries": [],
                "emergence_layer": None,
                "peak_layer": layers[0],
                "peak_probability": 0.0,
            }
        )

    for layer_idx in layers:
        h = hooks.state.hidden_states.get(layer_idx)
        if h is None:
            continue

        if h.ndim == 3:
            vec = h[0, token_position, :]
        elif h.ndim == 2:
            vec = h[token_position, :]
        else:
            vec = h

        pos_logits = _norm_project(final_norm, lm_head, vec)
        probs = mx.softmax(pos_logits, axis=-1)

        for i, (tid, _) in enumerate(resolved):
            prob = float(probs[tid])
            rank = int(mx.sum(probs > probs[tid]))
            is_top1 = rank == 0

            if is_top1 and trajectories[i]["emergence_layer"] is None:
                trajectories[i]["emergence_layer"] = layer_idx

            if prob > trajectories[i]["peak_probability"]:
                trajectories[i]["peak_probability"] = prob
                trajectories[i]["peak_layer"] = layer_idx

            trajectories[i]["entries"].append(
                CandidateLayerEntry(
                    layer=layer_idx,
                    probability=round(prob, 6),
                    rank=rank,
                    is_top1=is_top1,
                )
            )

    # Detect crossings
    crossings: list[CrossingEvent] = []
    prev_leader: str | None = None
    for layer_idx in layers:
        # Find who is top-1 at this layer
        best_prob = -1.0
        best_tok = ""
        for traj in trajectories:
            for entry in traj["entries"]:
                if entry.layer == layer_idx and entry.probability > best_prob:
                    best_prob = entry.probability
                    best_tok = traj["token"]

        if prev_leader is not None and best_tok != prev_leader:
            crossings.append(
                CrossingEvent(
                    layer=layer_idx,
                    previous_leader=prev_leader,
                    new_leader=best_tok,
                    new_leader_probability=round(best_prob, 6),
                )
            )
        prev_leader = best_tok

    # Build candidate trajectories
    candidate_results = []
    for traj in trajectories:
        candidate_results.append(
            CandidateTrajectory(
                token=traj["token"],
                token_id=traj["token_id"],
                layers=traj["entries"],
                emergence_layer=traj["emergence_layer"],
                peak_layer=traj["peak_layer"],
                peak_probability=round(traj["peak_probability"], 6),
            )
        )

    # Final winner
    final_winner = ""
    final_prob = 0.0
    for traj in trajectories:
        if traj["entries"] and traj["entries"][-1].probability > final_prob:
            final_prob = traj["entries"][-1].probability
            final_winner = traj["token"]

    result = TrackRaceResult(
        prompt=prompt,
        token_position=token_position,
        token_text=token_text,
        num_candidates=len(candidates),
        candidates=candidate_results,
        crossings=crossings,
        final_winner=final_winner,
        final_probability=round(final_prob, 6),
    )
    return result.model_dump(exclude_none=True)
