# tests/test_brain.py
import gzip
import json
from pathlib import Path

from fractus_vorax.brain import Brain
from fractus_vorax.compiler.atoms import Atom

QA = [
    Atom("what is the capital of france", "paris", "cap.csv:2"),
    Atom("what is the capital of spain", "madrid", "cap.csv:3"),
    Atom("what is the capital of japan", "tokyo", "cap.csv:4"),
    Atom("who wrote hamlet", "william shakespeare", "books.csv:2"),
]


def test_ingest_and_ask_top1():
    brain = Brain(D=2048)
    brain.ingest_atoms(QA)
    assert len(brain) == 4
    hits = brain.ask("what is the capital of france", k=2)
    assert hits[0][0].context == "paris"
    assert hits[0][1] > 0.3


def test_ask_generalizes_token_overlap():
    """Jamais vu mot à mot, mais les tokens partagés doivent hisser le bon domaine."""
    brain = Brain(D=2048)
    brain.ingest_atoms(QA)
    hits = brain.ask("what is the capital of germany", k=1)
    assert hits[0][0].statement.startswith("what is the capital of")


def test_save_load_roundtrip(tmp_path):
    brain = Brain(D=2048)
    brain.ingest_atoms(QA)
    brain.save(tmp_path / "brain")
    loaded = Brain.load(tmp_path / "brain")
    assert len(loaded) == 4
    assert loaded.ask("who wrote hamlet", k=1)[0][0].context == "william shakespeare"


def test_save_is_deterministic(tmp_path):
    b1, b2 = Brain(D=2048), Brain(D=2048)
    b1.ingest_atoms(QA)
    b2.ingest_atoms(QA)
    b1.save(tmp_path / "b1")
    b2.save(tmp_path / "b2")
    assert (tmp_path / "b1" / "meta.json").read_bytes() == (tmp_path / "b2" / "meta.json").read_bytes()
    assert (tmp_path / "b1" / "atoms.jsonl.gz").read_bytes() == (tmp_path / "b2" / "atoms.jsonl.gz").read_bytes()
    assert (tmp_path / "b1" / "traces.npz").read_bytes() == (tmp_path / "b2" / "traces.npz").read_bytes()