"""
inference/generate_stream.py — Real-time streaming generation for MiniTransformer.

Yields incremental text tokens as they are sampled autoregressively,
enabling responsive ChatGPT-style typewriter streaming in web interfaces.
"""

from typing import Generator
import torch
import torch.nn.functional as F

from tokenizer.bpe_tokenizer import BPETokenizer
from inference.generate import sample_logits


def generate_stream(
    model,
    prompt:             str,
    tokenizer:          "BPETokenizer | None" = None,
    max_new_tokens:     int   = 100,
    temperature:        float = 0.6,
    top_k:              int   = 40,
    top_p:              float = 0.9,
    repetition_penalty: float = 1.15,
    penalty_window:     int   = 64,
    device:             str   = "cuda" if torch.cuda.is_available() else "cpu",
) -> Generator[str, None, None]:
    """
    Generate text from a prompt, yielding the updated assistant response after each token.

    Yields:
        Full decoded assistant response text accumulated so far.
    """
    if tokenizer is None:
        tokenizer = BPETokenizer("dataProcessed/tokenizer.json")

    context_len = model.config.context_len
    model.eval()
    model.to(device)

    # Encode prompt
    prompt_ids = tokenizer.encode(prompt)
    tokens = torch.tensor(prompt_ids, dtype=torch.long, device=device).unsqueeze(0)
    prompt_len = len(prompt_ids)

    past_kvs = None
    temperature = max(temperature, 1e-5)
    generated_ids = []

    with torch.no_grad():
        for _ in range(max_new_tokens):
            if past_kvs is not None:
                context = tokens[:, -1:]
            else:
                context = tokens[:, -context_len:]

            logits, past_kvs_new = model(context, past_kvs)
            past_kvs = past_kvs_new

            next_logits = logits[:, -1, :].clone()

            # Repetition penalty
            if repetition_penalty > 1.0 and tokens.shape[1] > 0:
                recent = tokens[0, -penalty_window:].tolist()
                for tid in set(recent):
                    if next_logits[0, tid] > 0:
                        next_logits[0, tid] /= repetition_penalty
                    else:
                        next_logits[0, tid] *= repetition_penalty

            next_id = sample_logits(next_logits, temperature, top_k, top_p)

            if next_id == tokenizer.eos_id:
                break

            generated_ids.append(next_id)
            next_tensor = torch.tensor([[next_id]], dtype=torch.long, device=device)
            tokens = torch.cat([tokens, next_tensor], dim=1)

            # Robust BPE byte-level decoding
            full_decoded = tokenizer.decode(generated_ids)
            clean_text = tokenizer.clean(full_decoded)
            yield clean_text