"""Phase 2/3 — tests de la génération (échantillonnage écrit à la main)."""

from __future__ import annotations

import torch

from src.generate import _apply_top_k, _apply_top_p, generate, generate_text
from src.model.text_model import CortexTextLM, TextConfig
from src.tokenizer import load_tokenizer

TINY = TextConfig(vocab_size=32, n_layer=2, n_head=2, n_embd=32, ctx_len=16, dropout=0.0)


def _tiny():
    torch.manual_seed(0)
    return CortexTextLM(TINY).eval()


def test_top_k_keeps_exactly_k_finite():
    logits = torch.randn(1, 10)
    filtered = _apply_top_k(logits, 3)
    assert torch.isfinite(filtered).sum().item() == 3


def test_top_p_keeps_a_prefix_by_probability():
    logits = torch.zeros(1, 5)
    logits[0, 0] = 10.0  # proba ~ 1
    filtered = _apply_top_p(logits, 0.5)
    assert torch.isfinite(filtered[0, 0])
    assert not torch.isfinite(filtered[0, 1:]).any()


def test_greedy_is_deterministic():
    model = _tiny()
    x = torch.randint(0, 32, (1, 4))
    a = generate(model, x, max_new_tokens=5, greedy=True)
    b = generate(model, x, max_new_tokens=5, greedy=True)
    assert torch.equal(a, b)


def test_generate_extends_sequence():
    model = _tiny()
    x = torch.randint(0, 32, (1, 4))
    out = generate(model, x, max_new_tokens=6, temperature=1.0)
    assert out.shape == (1, 10)
    assert torch.equal(out[:, :4], x)


def test_generate_respects_ctx_len():
    model = _tiny()
    x = torch.randint(0, 32, (1, 16))
    out = generate(model, x, max_new_tokens=3, greedy=True)
    assert out.shape[1] == 19  # la condition est tronquée mais la sortie grandit


def test_generate_stops_on_eos(monkeypatch):
    model = _tiny()
    x = torch.randint(0, 32, (1, 3))
    eos = 7
    # Force le modèle à toujours prédire eos.
    def fake_forward(input_ids, targets=None):
        logits = torch.full((input_ids.size(0), input_ids.size(1), 32), -100.0)
        logits[:, :, eos] = 100.0
        return logits, None

    monkeypatch.setattr(model, "forward", fake_forward)
    out = generate(model, x, max_new_tokens=50, greedy=True, eos_id=eos)
    assert out.shape[1] == 4  # un seul nouveau token, puis arrêt


def test_generate_text_with_real_tokenizer():
    """Bout-en-bout avec le vrai tokenizer 32k : ne doit pas planter."""
    tok = load_tokenizer()
    model = CortexTextLM(
        TextConfig(vocab_size=tok.vocab_size, n_layer=1, n_head=2, n_embd=32, ctx_len=32)
    ).eval()
    text = generate_text(model, tok, "Bonjour", max_new_tokens=5, greedy=True)
    assert isinstance(text, str) and len(text) > 0