# bclm/inference.py
"""
BCLM-1 Inference — Optimized [B, T, C] Layout

Efficient autoregressive generation with KV and oscillator state caching.
Supports Grouped Query Attention (GQA) and bigram hash embedding injection.

Provides:
  - BCLM1ForGeneration  — prefill + cached single-token decode
  - ChatSession          — multi-turn conversation with streaming
"""
from __future__ import annotations

import inspect
import sys
from dataclasses import dataclass
from functools import lru_cache
from typing import Dict, Iterator, List, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

from .config import BCLM1ModelConfig, EOT_TOKEN_ID
from .tokenizer import BCLMTokenizer, SPECIAL_TOKEN_IDS, ROLE_TO_TOKEN


# =============================================================================
# GQA Utilities
# =============================================================================


@lru_cache(maxsize=None)
def _sdpa_supports_enable_gqa() -> bool:
    try:
        sig = inspect.signature(F.scaled_dot_product_attention)
        return "enable_gqa" in sig.parameters
    except Exception:
        return False


def _repeat_kv_heads(x: torch.Tensor, n_heads: int) -> torch.Tensor:
    b, hkv, t, d = x.shape
    if hkv == n_heads:
        return x
    if n_heads % hkv != 0:
        raise ValueError(f"Cannot repeat kv heads: Hq={n_heads}, Hkv={hkv}")
    return x.repeat_interleave(n_heads // hkv, dim=1)


# =============================================================================
# Bigram Hash (mirrors model._compute_bigram_hash constants)
# =============================================================================

_BIGRAM_RAND_INT_1 = 36313
_BIGRAM_RAND_INT_2 = 27191


def _compute_bigram_hash_full(
    x: torch.Tensor,
    bigram_vocab_size: int,
) -> torch.Tensor:
    mod = bigram_vocab_size - 1
    result = torch.empty_like(x)
    result[:, 0] = mod
    result[:, 1:] = (
        torch.bitwise_xor(
            _BIGRAM_RAND_INT_1 * x[:, 1:],
            _BIGRAM_RAND_INT_2 * x[:, :-1],
        )
        % mod
    )
    return result


def _compute_bigram_hash_step(
    curr_token: torch.Tensor,
    prev_token: torch.Tensor,
    bigram_vocab_size: int,
) -> torch.Tensor:
    mod = bigram_vocab_size - 1
    return (
        torch.bitwise_xor(
            _BIGRAM_RAND_INT_1 * curr_token,
            _BIGRAM_RAND_INT_2 * prev_token,
        )
        % mod
    )


# =============================================================================
# Inference Cache
# =============================================================================


@dataclass
class LayerCache:
    osc_complex: torch.Tensor
    osc_real: torch.Tensor
    mixer_conv_cache: torch.Tensor
    k_cache: Optional[torch.Tensor] = None
    v_cache: Optional[torch.Tensor] = None


@dataclass
class BCLM1Cache:
    config: BCLM1ModelConfig
    input_conv_cache: torch.Tensor
    layers: List[LayerCache]
    local_attn_layer_set: set
    global_attn_layer_set: set
    seq_pos: int = 0
    last_token_id: Optional[torch.Tensor] = None

    @classmethod
    def create(
        cls,
        config: BCLM1ModelConfig,
        batch_size: int,
        device: torch.device,
        dtype: torch.dtype = torch.bfloat16,
    ) -> "BCLM1Cache":
        dim = config.embed_dim

        input_conv_cache = torch.zeros(
            batch_size, dim, config.conv_kernel_size - 1,
            device=device, dtype=dtype,
        )

        local_attn_layer_set = set(int(x) for x in config.local_attn_layers)
        global_attn_layer_set = set(int(x) for x in config.global_attn_layers)

        overlap = local_attn_layer_set & global_attn_layer_set
        if overlap:
            raise ValueError(
                f"Layers cannot be both local and global attention: {sorted(overlap)}"
            )

        layers: List[LayerCache] = []
        for _ in range(config.n_layers):
            layers.append(
                LayerCache(
                    osc_complex=torch.zeros(
                        batch_size, dim, config.osc_n_pairs,
                        device=device, dtype=torch.complex64,
                    ),
                    osc_real=torch.zeros(
                        batch_size, dim, config.osc_n_real,
                        device=device, dtype=dtype,
                    ),
                    mixer_conv_cache=torch.zeros(
                        batch_size, dim, 1,
                        device=device, dtype=dtype,
                    ),
                )
            )

        return cls(
            config=config,
            input_conv_cache=input_conv_cache,
            layers=layers,
            local_attn_layer_set=local_attn_layer_set,
            global_attn_layer_set=global_attn_layer_set,
            seq_pos=0,
            last_token_id=None,
        )

    def clear(self) -> None:
        self.input_conv_cache.zero_()
        for layer in self.layers:
            layer.osc_complex.zero_()
            layer.osc_real.zero_()
            layer.k_cache = None
            layer.v_cache = None
        self.seq_pos = 0
        self.last_token_id = None


# =============================================================================
# Cached Operations
# =============================================================================


def cached_conv_step(
    weight: torch.Tensor,
    x_ct: torch.Tensor,
    cache: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    full_input = torch.cat([cache, x_ct], dim=-1)
    out = torch.einsum("oik,bik->bo", weight, full_input).unsqueeze(-1)
    new_cache = full_input[:, :, 1:].clone()
    return out, new_cache


def cached_oscillator_step(
    x: torch.Tensor,
    complex_state: torch.Tensor,
    real_state: torch.Tensor,
    log_alpha: torch.Tensor,
    log_omega: torch.Tensor,
    phi: torch.Tensor,
    log_amp: torch.Tensor,
    log_beta: torch.Tensor,
    log_amp2: torch.Tensor,
    sign2: torch.Tensor,
    clamp_min_decay: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    x_squeezed = x.squeeze(1).float()

    alpha = torch.exp(log_alpha).clamp_min(clamp_min_decay)
    omega = torch.exp(log_omega)
    z = torch.exp((-alpha + 1j * omega).to(torch.complex64))

    new_complex = (
        x_squeezed.unsqueeze(-1).to(torch.complex64)
        + z.unsqueeze(0) * complex_state
    )

    amp = torch.exp(log_amp)
    phase = torch.exp(1j * phi.to(torch.complex64))
    complex_out = (
        amp.unsqueeze(0) * (phase.unsqueeze(0) * new_complex).real
    ).sum(dim=-1)

    beta = torch.exp(log_beta).clamp_min(clamp_min_decay)
    d = torch.exp(-beta)

    new_real = x_squeezed.unsqueeze(-1) + d.unsqueeze(0) * real_state

    amp2 = torch.exp(log_amp2)
    s2 = torch.tanh(sign2)
    real_out = ((s2 * amp2).unsqueeze(0) * new_real).sum(dim=-1)

    out = (complex_out + real_out).unsqueeze(1).to(x.dtype)
    return out, new_complex, new_real


def cached_attention_step(
    x: torch.Tensor,
    k_cache: Optional[torch.Tensor],
    v_cache: Optional[torch.Tensor],
    norm_weight: Optional[torch.Tensor],
    qkv_weight: torch.Tensor,
    out_proj_weight: torch.Tensor,
    n_heads: int,
    kv_heads: int,
    head_dim: int,
    max_cache_len: Optional[int],
    norm_eps: float = 1e-8,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    B = x.size(0)
    dim = x.size(-1)
    kv_dim = kv_heads * head_dim

    x_normed = F.rms_norm(
        x.squeeze(1), (dim,), weight=norm_weight, eps=norm_eps,
    )

    qkv = F.linear(x_normed, qkv_weight)
    q, k, v = qkv.split([dim, kv_dim, kv_dim], dim=-1)

    q = q.view(B, 1, n_heads, head_dim)
    k = k.view(B, 1, kv_heads, head_dim)
    v = v.view(B, 1, kv_heads, head_dim)

    if k_cache is None:
        k_full = k
        v_full = v
    else:
        k_full = torch.cat([k_cache, k], dim=1)
        v_full = torch.cat([v_cache, v], dim=1)

    if max_cache_len is not None and k_full.size(1) > max_cache_len:
        k_full = k_full[:, -max_cache_len:]
        v_full = v_full[:, -max_cache_len:]

    qh = q.transpose(1, 2)
    kh = k_full.transpose(1, 2)
    vh = v_full.transpose(1, 2)

    if kv_heads != n_heads:
        if _sdpa_supports_enable_gqa():
            out = F.scaled_dot_product_attention(
                qh, kh, vh,
                attn_mask=None, dropout_p=0.0,
                is_causal=False, enable_gqa=True,
            )
        else:
            kh_exp = _repeat_kv_heads(kh, n_heads)
            vh_exp = _repeat_kv_heads(vh, n_heads)
            out = F.scaled_dot_product_attention(
                qh, kh_exp, vh_exp,
                attn_mask=None, dropout_p=0.0, is_causal=False,
            )
    else:
        out = F.scaled_dot_product_attention(
            qh, kh, vh,
            attn_mask=None, dropout_p=0.0, is_causal=False,
        )

    out = out.transpose(1, 2).contiguous().view(B, dim)
    out = F.linear(out, out_proj_weight)

    return out.unsqueeze(1), k_full, v_full


# =============================================================================
# Generation Wrapper
# =============================================================================


class BCLM1ForGeneration(nn.Module):
    """
    BCLM-1 wrapper optimised for autoregressive generation.

    Provides prefill + cached single-token decode, batch generation,
    and streaming generation.
    """

    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model
        self.config = model.config

    # ── Prefill ─────────────────────────────────────────────

    @torch.no_grad()
    def prefill(
        self,
        idx: torch.Tensor,
        cache: BCLM1Cache,
    ) -> torch.Tensor:
        B, T = idx.shape
        model = self.model
        config = self.config

        tok = model.tok_emb(idx)
        x = tok + model.input_mixer(tok)

        ks = config.conv_kernel_size
        tok_ct = tok.transpose(1, 2)
        if T >= ks - 1:
            cache.input_conv_cache = tok_ct[:, :, -(ks - 1):].clone()
        else:
            cache.input_conv_cache.zero_()
            cache.input_conv_cache[:, :, -T:] = tok_ct

        bigram_vocab_size = model._bigram_vocab_size
        bigram_idx = _compute_bigram_hash_full(idx, bigram_vocab_size)
        x0_bigram = model.bigram_embed(bigram_idx)

        cache.last_token_id = idx[:, -1:].clone()

        local_set = cache.local_attn_layer_set
        global_set = cache.global_attn_layer_set

        for layer_idx, layer in enumerate(model.layers):
            x = x + model.bigram_lambdas[layer_idx] * x0_bigram

            residual = x
            x_norm = layer.mixer.norm(x)
            u_raw = layer.mixer.in_proj(x_norm)
            u, z = u_raw.chunk(2, dim=-1)

            u_filt = layer.mixer.osc(u)

            lc = cache.layers[layer_idx]
            lc.osc_complex, lc.osc_real = self._compute_final_osc_state(
                u, layer.mixer.osc,
            )

            out = u_filt * (z + layer.mixer.gate_bias)
            x = layer.mixer.norm(residual + out)

            if (layer_idx in local_set) or (layer_idx in global_set):
                attn = layer.attn
                if (
                    hasattr(attn, "qkv")
                    and hasattr(attn, "norm")
                    and hasattr(attn, "out_proj")
                ):
                    n_heads = attn.n_heads
                    kv_heads = attn.kv_heads
                    head_dim = attn.head_dim
                    kv_dim = kv_heads * head_dim

                    x_norm_a = attn.norm(x)
                    qkv = attn.qkv(x_norm_a)
                    _, k, v = qkv.split([attn.dim, kv_dim, kv_dim], dim=-1)

                    k = k.view(B, T, kv_heads, head_dim)
                    v = v.view(B, T, kv_heads, head_dim)

                    if layer_idx in local_set:
                        ws = int(
                            getattr(attn, "window_size", config.attn_window_size)
                        )
                        trim = ws
                    else:
                        trim = int(config.max_seq_len)

                    if T > trim:
                        lc.k_cache = k[:, -trim:].clone()
                        lc.v_cache = v[:, -trim:].clone()
                    else:
                        lc.k_cache = k.clone()
                        lc.v_cache = v.clone()

                    x = attn(x)

            residual = x
            x_norm_f = layer.ffn.norm(x)
            v_ffn, g = layer.ffn.proj_in(x_norm_f).chunk(2, dim=-1)
            x_out = layer.ffn.proj_out(F.gelu(g, approximate="tanh") * v_ffn)
            x = residual + x_out

        x = model.norm_f(x)
        logits = model.head(x)

        cache.seq_pos = T
        return logits[:, -1, :]

    # ── Oscillator State Seeding ────────────────────────────

    def _compute_final_osc_state(
        self,
        u: torch.Tensor,
        osc,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        B, T, dim = u.shape
        device = u.device

        complex_state = torch.zeros(
            B, dim, osc.log_alpha.size(1),
            device=device, dtype=torch.complex64,
        )
        real_state = torch.zeros(
            B, dim, osc.log_beta.size(1),
            device=device, dtype=u.dtype,
        )

        alpha = torch.exp(osc.log_alpha).clamp_min(osc.clamp_min_decay)
        omega = torch.exp(osc.log_omega)
        z = torch.exp((-alpha + 1j * omega).to(torch.complex64))

        beta = torch.exp(osc.log_beta).clamp_min(osc.clamp_min_decay)
        d = torch.exp(-beta)

        for t in range(T):
            x_t = u[:, t, :].float()
            complex_state = (
                x_t.unsqueeze(-1).to(torch.complex64)
                + z.unsqueeze(0) * complex_state
            )
            real_state = x_t.unsqueeze(-1) + d.unsqueeze(0) * real_state

        return complex_state, real_state

    # ── Single-Token Decode ─────────────────────────────────

    @torch.no_grad()
    def decode_one_token(
        self,
        idx: torch.Tensor,
        cache: BCLM1Cache,
    ) -> torch.Tensor:
        model = self.model
        config = self.config
        local_set = cache.local_attn_layer_set
        global_set = cache.global_attn_layer_set

        tok = model.tok_emb(idx)

        tok_ct = tok.transpose(1, 2)
        conv_out_ct, cache.input_conv_cache = cached_conv_step(
            model.input_mixer.conv.weight,
            tok_ct,
            cache.input_conv_cache,
        )
        x = tok + conv_out_ct.transpose(1, 2)

        bigram_vocab_size = model._bigram_vocab_size
        if cache.last_token_id is not None:
            bigram_hash = _compute_bigram_hash_step(
                idx, cache.last_token_id, bigram_vocab_size,
            )
        else:
            bigram_hash = torch.full_like(idx, bigram_vocab_size - 1)
        x0_bigram = model.bigram_embed(bigram_hash)

        cache.last_token_id = idx.clone()

        for layer_idx, layer in enumerate(model.layers):
            lc = cache.layers[layer_idx]

            x = x + model.bigram_lambdas[layer_idx] * x0_bigram

            residual = x
            x_norm = layer.mixer.norm(x)
            u_raw = layer.mixer.in_proj(x_norm)
            u, z = u_raw.chunk(2, dim=-1)

            osc = layer.mixer.osc
            u_filt, lc.osc_complex, lc.osc_real = cached_oscillator_step(
                u,
                lc.osc_complex,
                lc.osc_real,
                osc.log_alpha,
                osc.log_omega,
                osc.phi,
                osc.log_amp,
                osc.log_beta,
                osc.log_amp2,
                osc.sign2,
                osc.clamp_min_decay,
            )

            out = u_filt * (z + layer.mixer.gate_bias)
            x = layer.mixer.norm(residual + out)

            if (layer_idx in local_set) or (layer_idx in global_set):
                attn = layer.attn
                if (
                    hasattr(attn, "qkv")
                    and hasattr(attn, "norm")
                    and hasattr(attn, "out_proj")
                ):
                    n_heads = attn.n_heads
                    kv_heads = attn.kv_heads
                    head_dim = attn.head_dim

                    if layer_idx in local_set:
                        max_cache_len = int(
                            getattr(attn, "window_size", config.attn_window_size)
                        )
                    else:
                        max_cache_len = int(config.max_seq_len)

                    attn_out, lc.k_cache, lc.v_cache = cached_attention_step(
                        x,
                        lc.k_cache,
                        lc.v_cache,
                        attn.norm.weight,
                        attn.qkv.weight,
                        attn.out_proj.weight,
                        n_heads,
                        kv_heads,
                        head_dim,
                        max_cache_len,
                    )
                    x = x + attn_out

            residual = x
            x_norm_f = layer.ffn.norm(x)
            v_ffn, g = layer.ffn.proj_in(x_norm_f).chunk(2, dim=-1)
            x_out = layer.ffn.proj_out(F.gelu(g, approximate="tanh") * v_ffn)
            x = residual + x_out

        x = model.norm_f(x)
        logits = model.head(x)

        cache.seq_pos += 1
        return logits[:, 0, :]

    # ── Batch Generation ────────────────────────────────────

    @torch.no_grad()
    def generate(
        self,
        prompt_tokens: torch.Tensor,
        max_new_tokens: int,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        stop_token: Optional[int] = EOT_TOKEN_ID,
    ) -> torch.Tensor:
        """
        Autoregressive generation with prefill + cached decode.

        Returns the full sequence including prompt: [B, T + generated].
        """
        B = prompt_tokens.size(0)
        device = prompt_tokens.device
        dtype = next(self.model.parameters()).dtype

        cache = BCLM1Cache.create(self.config, B, device, dtype)

        logits = self.prefill(prompt_tokens, cache)
        generated = prompt_tokens.clone()

        for _ in range(max_new_tokens):
            next_token = self._sample(logits, temperature, top_k, top_p)
            generated = torch.cat([generated, next_token.unsqueeze(1)], dim=1)

            if stop_token is not None and (next_token == stop_token).all():
                break

            logits = self.decode_one_token(next_token.unsqueeze(1), cache)

        return generated

    # ── Streaming Generation ────────────────────────────────

    @torch.no_grad()
    def generate_stream(
        self,
        prompt_tokens: torch.Tensor,
        max_new_tokens: int,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        stop_token: Optional[int] = EOT_TOKEN_ID,
    ) -> Iterator[int]:
        """
        Streaming autoregressive generation — yields one token ID at a
        time. Only supports batch_size=1.

        The stop token is yielded (so callers can distinguish natural
        stops from length truncation).
        """
        B = prompt_tokens.size(0)
        if B != 1:
            raise ValueError(f"generate_stream requires batch_size=1, got {B}")

        device = prompt_tokens.device
        dtype = next(self.model.parameters()).dtype

        cache = BCLM1Cache.create(self.config, B, device, dtype)
        logits = self.prefill(prompt_tokens, cache)

        for _ in range(max_new_tokens):
            next_token = self._sample(logits, temperature, top_k, top_p)
            token_id = int(next_token.item())
            yield token_id

            if stop_token is not None and token_id == stop_token:
                return

            logits = self.decode_one_token(next_token.unsqueeze(1), cache)

    # ── Sampling ────────────────────────────────────────────

    @staticmethod
    def _sample(
        logits: torch.Tensor,
        temperature: float,
        top_k: Optional[int],
        top_p: Optional[float],
    ) -> torch.Tensor:
        if temperature <= 0:
            return logits.argmax(dim=-1)

        logits = logits / temperature

        if top_k is not None and top_k > 0:
            top_k = min(top_k, logits.size(-1))
            indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
            logits = logits.masked_fill(indices_to_remove, float("-inf"))

        if top_p is not None and top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            cumulative_probs = torch.cumsum(
                F.softmax(sorted_logits, dim=-1), dim=-1,
            )

            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
                ..., :-1
            ].clone()
            sorted_indices_to_remove[..., 0] = 0

            indices_to_remove = sorted_indices_to_remove.scatter(
                -1, sorted_indices, sorted_indices_to_remove,
            )
            logits = logits.masked_fill(indices_to_remove, float("-inf"))

        probs = F.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1).squeeze(-1)


# =============================================================================
# Chat Session
# =============================================================================


_UNSET = object()


class ChatSession:
    """
    Multi-turn conversation session with streaming support.

    Manages conversation history, prompt encoding, and generation
    parameter defaults. Suitable for both programmatic use and
    interactive CLI sessions.
    """

    def __init__(
        self,
        gen_model: BCLM1ForGeneration,
        tokenizer: BCLMTokenizer,
        *,
        system_prompt: Optional[str] = None,
        max_new_tokens: int = 512,
        temperature: float = 1.0,
        top_k: Optional[int] = 50,
        top_p: Optional[float] = 0.9,
    ) -> None:
        self.gen_model = gen_model
        self.tokenizer = tokenizer
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.top_k = top_k
        self.top_p = top_p

        self._messages: List[Dict[str, str]] = []
        self._system_prompt = system_prompt
        if system_prompt:
            self._messages.append(
                {"role": "system", "content": system_prompt}
            )

    @property
    def messages(self) -> List[Dict[str, str]]:
        """A copy of the conversation history."""
        return list(self._messages)

    @property
    def device(self) -> torch.device:
        return next(self.gen_model.model.parameters()).device

    def clear(self) -> None:
        """Reset conversation, keeping the system prompt if one was set."""
        self._messages.clear()
        if self._system_prompt:
            self._messages.append(
                {"role": "system", "content": self._system_prompt}
            )

    def _resolve(self, val, default):
        return default if val is _UNSET else val

    def _build_prompt_tensor(self, reserve_tokens: int) -> torch.Tensor:
        ids = self.tokenizer.encode_chat(
            self._messages, add_generation_prompt=True,
        )

        max_prompt = self.gen_model.config.max_seq_len - reserve_tokens
        if max_prompt < 1:
            max_prompt = 1

        if len(ids) > max_prompt:
            ids = ids[-max_prompt:]

        return torch.tensor([ids], dtype=torch.long, device=self.device)

    # ── send() — full response ──────────────────────────────

    def send(
        self,
        user_message: str,
        *,
        max_new_tokens=_UNSET,
        temperature=_UNSET,
        top_k=_UNSET,
        top_p=_UNSET,
    ) -> str:
        """
        Send a user message and return the complete assistant response.

        Per-call overrides fall back to session defaults when omitted.
        """
        _max = self._resolve(max_new_tokens, self.max_new_tokens)
        _temp = self._resolve(temperature, self.temperature)
        _topk = self._resolve(top_k, self.top_k)
        _topp = self._resolve(top_p, self.top_p)

        self._messages.append({"role": "user", "content": user_message})
        prompt = self._build_prompt_tensor(reserve_tokens=_max)

        response_ids: List[int] = []
        for tid in self.gen_model.generate_stream(
            prompt, _max,
            temperature=_temp, top_k=_topk, top_p=_topp,
            stop_token=EOT_TOKEN_ID,
        ):
            if tid == EOT_TOKEN_ID:
                break
            response_ids.append(tid)

        response_text = self.tokenizer.decode_response(response_ids)
        self._messages.append(
            {"role": "assistant", "content": response_text}
        )
        return response_text

    # ── send_stream() — yields text chunks ──────────────────

    def send_stream(
        self,
        user_message: str,
        *,
        max_new_tokens=_UNSET,
        temperature=_UNSET,
        top_k=_UNSET,
        top_p=_UNSET,
    ) -> Iterator[str]:
        """
        Send a user message and yield the assistant response as
        incremental text chunks suitable for real-time display.

        History is updated after the full response has been yielded.
        """
        _max = self._resolve(max_new_tokens, self.max_new_tokens)
        _temp = self._resolve(temperature, self.temperature)
        _topk = self._resolve(top_k, self.top_k)
        _topp = self._resolve(top_p, self.top_p)

        self._messages.append({"role": "user", "content": user_message})
        prompt = self._build_prompt_tensor(reserve_tokens=_max)

        response_ids: List[int] = []
        prev_text = ""

        for tid in self.gen_model.generate_stream(
            prompt, _max,
            temperature=_temp, top_k=_topk, top_p=_topp,
            stop_token=EOT_TOKEN_ID,
        ):
            if tid == EOT_TOKEN_ID:
                break
            if tid in SPECIAL_TOKEN_IDS:
                continue

            response_ids.append(tid)

            try:
                full_text = self.tokenizer.decode(list(response_ids))
            except Exception:
                continue

            delta = full_text[len(prev_text):]
            if delta:
                prev_text = full_text
                yield delta

        # Final flush
        if response_ids:
            try:
                final = self.tokenizer.decode(response_ids)
                delta = final[len(prev_text):]
                if delta:
                    yield delta
            except Exception:
                pass

        response_text = (
            self.tokenizer.decode_response(response_ids)
            if response_ids else ""
        )
        self._messages.append(
            {"role": "assistant", "content": response_text}
        )

    # ── interactive() — CLI chat loop ───────────────────────

    def interactive(self) -> None:
        """
        Blocking interactive chat loop.

        Commands:
          /clear    — reset conversation
          /history  — print message list
          /help     — show commands
          Ctrl-C    — exit
        """
        print("BCLM Chat  (type /help for commands, Ctrl-C to exit)")
        print("─" * 56)
        if self._system_prompt:
            print(f"System: {self._system_prompt}")
            print("─" * 56)

        while True:
            try:
                user_input = input("\nYou: ").strip()
            except (KeyboardInterrupt, EOFError):
                print("\nGoodbye.")
                return

            if not user_input:
                continue

            if user_input.startswith("/"):
                cmd = user_input.lower().split()[0]
                if cmd == "/clear":
                    self.clear()
                    print("[chat] Conversation cleared.")
                    continue
                if cmd == "/history":
                    for i, m in enumerate(self._messages):
                        print(f"  [{i}] {m['role']}: {m['content'][:80]}...")
                    continue
                if cmd == "/help":
                    print("  /clear    — reset conversation")
                    print("  /history  — show message list")
                    print("  /help     — this message")
                    print("  Ctrl-C    — exit")
                    continue
                print(f"[chat] Unknown command: {cmd}  (try /help)")
                continue

            sys.stdout.write("\nAssistant: ")
            sys.stdout.flush()
            for chunk in self.send_stream(user_input):
                sys.stdout.write(chunk)
                sys.stdout.flush()
            print()