"""Tests for conversational Q/A training and respond() (``palimseste.lm``).

Verifies:
  - train_on_qa_pairs writes traces to M (O(1) per token)
  - respond() returns a non-empty string after training
  - respond() recovers the exact trained answer for a seen question (greedy)
  - respond() handles unseen questions gracefully (no crash)
  - the episode format (BOS q EOS BOS a EOS) is consistent
  - generate_from_ids respects stop_on_eos
"""

from __future__ import annotations

import numpy as np
import pytest

from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
from palimseste.tokenizer import BOS, EOS


def _lm(D=5000, radius=30, ctx_window=48):
    cfg = PalimpsesteConfig(D=D, context_window=ctx_window, kernel_radius=radius,
                            temperature=0.3)
    return PalimpsesteForCausalLM(config=cfg)


# ----------------------------------------------------------------- train
def test_train_on_qa_pairs_grows_memory():
    lm = _lm()
    pairs = [("hello", "hi there"), ("bye", "goodbye")]
    full = "".join(q + a for q, a in pairs)
    lm.build_tokenizer(full)
    n = lm.train_on_qa_pairs(pairs)
    assert n > 0
    assert len(lm.mem) == n


def test_train_on_qa_pairs_requires_tokenizer():
    lm = _lm()
    with pytest.raises(AssertionError):
        lm.train_on_qa_pairs([("q", "a")])


# ----------------------------------------------------------------- respond
def test_respond_returns_exact_answer_for_seen_question():
    lm = _lm(radius=0)
    pairs = [
        ("bonjour", "salut je suis palimpseste"),
        ("qui es-tu", "je suis palimpseste"),
        ("comment tu apprends", "j'apprends en ecrivant"),
    ]
    full = "".join(q + a for q, a in pairs)
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    # greedy: should recover the exact trained answer
    resp = lm.respond("bonjour", temperature=0.0, seed=0)
    assert resp == "salut je suis palimpseste"


def test_respond_recovers_multiple_answers():
    lm = _lm(radius=0)
    pairs = [
        ("hello", "hi i am palimpseste"),
        ("who are you", "i am palimpseste a hypervectorial cortex"),
        ("how do you learn", "i learn by writing to memory"),
        ("do you use a gpu", "no i do not use a gpu"),
    ]
    full = "".join(q + a for q, a in pairs)
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    for q, expected in pairs:
        resp = lm.respond(q, temperature=0.0, seed=0)
        assert resp == expected, f"Q={q!r}: got {resp!r}, expected {expected!r}"


def test_respond_unseen_question_no_crash():
    lm = _lm()
    pairs = [("hello", "hi there")]
    full = "hellohi there"
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    # unseen question should not crash; may return empty or partial
    resp = lm.respond("what is the weather", temperature=0.5, seed=0)
    assert isinstance(resp, str)


def test_respond_empty_input_no_crash():
    lm = _lm()
    lm.build_tokenizer("abc")
    lm.train_on_qa_pairs([("a", "b")])
    resp = lm.respond("", temperature=0.0, seed=0)
    assert isinstance(resp, str)


def test_respond_respects_max_tokens():
    lm = _lm(radius=0)
    # long answer
    pairs = [("q", "a" * 100)]
    full = "q" + "a" * 100
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    resp = lm.respond("q", max_new_tokens=5, temperature=0.0, seed=0)
    assert len(resp) <= 5


# ----------------------------------------------------------------- generate_from_ids
def test_generate_from_ids_stop_on_eos():
    lm = _lm(radius=0)
    text = "hello world"
    lm.build_tokenizer(text)
    lm.train_on_text(text)
    ids = lm.tokenizer.encode("hello", add_bos=True)
    out = lm.generate_from_ids(ids, max_new_tokens=50, temperature=0.0, seed=0,
                               stop_on_eos=True)
    assert isinstance(out, str)


def test_generate_from_ids_cold_start():
    lm = _lm()
    lm.build_tokenizer("abc")
    # untrained context -> cold start -> empty string
    out = lm.generate_from_ids([BOS], max_new_tokens=10, temperature=0.0, seed=0)
    assert isinstance(out, str)


# ----------------------------------------------------------------- episode format
def test_qa_episode_format_consistent():
    """The episode BOS q EOS BOS a EOS must be used consistently in
    train_on_qa_pairs and respond()."""
    lm = _lm(radius=0)
    pairs = [("hi", "hello there")]
    full = "hihello there"
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    # respond() must encode the question the same way: BOS q EOS BOS
    # then generate. The first generated token should be 'h' (start of "hello there")
    q_ids = lm.tokenizer.encode("hi", add_bos=True, add_eos=True)
    ctx = q_ids + [BOS]
    logits = lm._logits(ctx)
    import numpy as np
    top = int(np.argmax(logits))
    # the first answer char is 'h'
    assert lm.tokenizer.id2char[top] == "h"


# ----------------------------------------------------------------- robustness
def test_qa_with_french_accents():
    lm = _lm(radius=0)
    pairs = [("comment ça va", "ça va bien merci")]
    full = "comment ça vaça va bien merci"
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    resp = lm.respond("comment ça va", temperature=0.0, seed=0)
    assert "ça va bien" in resp or resp == "ça va bien merci"


def test_qa_multiple_pairs_no_interference():
    """Training on multiple Q/A pairs should not cause one answer to bleed
    into another (each episode is independent with its own BOS/EOS boundary)."""
    lm = _lm(radius=0)
    pairs = [
        ("aaa", "xxx"),
        ("bbb", "yyy"),
        ("ccc", "zzz"),
    ]
    full = "".join(q + a for q, a in pairs)
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    assert lm.respond("aaa", temperature=0.0, seed=0) == "xxx"
    assert lm.respond("bbb", temperature=0.0, seed=0) == "yyy"
    assert lm.respond("ccc", temperature=0.0, seed=0) == "zzz"