import json
import tempfile
from pathlib import Path

from markovian_rsa_mlx.audit import (
    AuditWriter,
    RunStartEvent,
    TraceCompleteEvent,
    TailExtractedEvent,
    AggregationPromptEvent,
    RoundCompleteEvent,
    FinalEvent,
    RunEndEvent,
)
from markovian_rsa_mlx.config import RSAConfig


def test_writer_emits_jsonl(tmp_path: Path):
    audit_path = tmp_path / "audit.jsonl"
    writer = AuditWriter(audit_path)
    writer.write(RunStartEvent(
        run_id="run-1",
        model_id="kyr0/zaya1-base-8b-MLX",
        quantization="q4_g64",
        config=RSAConfig(),
        prompt="2+2",
        created_at="2026-05-10T00:00:00Z",
    ))
    writer.write(TraceCompleteEvent(
        run_id="run-1",
        round=0,
        trace_id="t0",
        text="The answer is 4.",
        token_ids=[1, 2, 3],
        generated_tokens=3,
        finish_reason="eos",
        elapsed_s=0.5,
    ))
    writer.close()

    lines = audit_path.read_text().strip().split("\n")
    assert len(lines) == 2
    e0 = json.loads(lines[0])
    assert e0["event"] == "run_start"
    assert e0["model_id"] == "kyr0/zaya1-base-8b-MLX"
    assert e0["config"]["rounds"] == 2
    e1 = json.loads(lines[1])
    assert e1["event"] == "trace_complete"
    assert e1["text"] == "The answer is 4."
    assert e1["finish_reason"] == "eos"


def test_event_serialization_includes_event_field():
    ev = TailExtractedEvent(
        run_id="r", round=0, trace_id="t",
        tail_token_ids=[1, 2, 3], tail_text="hi", tail_tokens=3,
    )
    payload = ev.to_dict()
    assert payload["event"] == "tail_extracted"


def test_writer_appends_subsequent_events(tmp_path: Path):
    audit_path = tmp_path / "audit.jsonl"
    w = AuditWriter(audit_path)
    w.write(RunEndEvent(run_id="r", elapsed_s=1.0,
                        total_generated_tokens=10, peak_memory_bytes=1000))
    w.close()
    w2 = AuditWriter(audit_path, mode="a")
    w2.write(RunEndEvent(run_id="r2", elapsed_s=2.0,
                         total_generated_tokens=20, peak_memory_bytes=2000))
    w2.close()
    lines = audit_path.read_text().strip().split("\n")
    assert len(lines) == 2


def test_writer_no_op_when_path_none():
    w = AuditWriter(None)
    w.write(RunEndEvent(run_id="x", elapsed_s=1.0,
                        total_generated_tokens=0, peak_memory_bytes=0))
    w.close()  # no exception