"""Tests for bandito/assertions.py."""

from __future__ import annotations

import json
import os
from typing import Any, Iterator

from bandito.assertions import AssertionRunResult, assert_traces
from bandito.models import Trace
from bandito.storage.base import StoreBase


class _MockStore(StoreBase):
    """In-memory store for testing."""

    def __init__(self, traces: list[dict]):
        self._traces = traces

    def write_trace(self, trace: Trace) -> None:
        self._traces.append(trace.to_dict())

    def _read_trace(self, trace_id: str) -> dict[str, Any] | None:
        for t in self._traces:
            if t.get("trace_id") == trace_id:
                return t
        return None

    def iter_traces(self, project: str | None = None) -> Iterator[dict[str, Any]]:
        for t in self._traces:
            if project is not None and t.get("project") != project:
                continue
            yield t


def _make_traces(n: int, project: str = "bot") -> list[dict]:
    return [
        {
            "trace_id": f"tr-{i}",
            "project": project,
            "input": f"question {i}",
            "output": f"answer {i}" if i % 3 != 0 else "",
            "started_at": f"2026-03-20T12:0{i}:00Z",
            "status": "ok",
            "spans": [
                {
                    "span_id": f"sp-{i}",
                    "kind": "llm",
                    "cost": 0.01 * (i + 1),
                    "status": "error" if i == 4 else "ok",
                }
            ],
        }
        for i in range(n)
    ]


# -- Assertion functions for testing --


def output_not_empty(trace):
    return bool(trace.get("output"))


def cost_under_threshold(trace):
    total = sum(s.get("cost", 0) for s in trace.get("spans", []))
    return total < 0.04


def no_error_spans(trace):
    return all(s.get("status") != "error" for s in trace.get("spans", []))


def always_raises(trace):
    raise ValueError("intentional test error")


# -- Tests --


class TestAssertTraces:
    def test_all_pass(self):
        traces = [
            {"trace_id": "tr-0", "project": "bot", "output": "hi", "spans": []},
            {"trace_id": "tr-1", "project": "bot", "output": "there", "spans": []},
        ]
        store = _MockStore(traces)

        result = assert_traces("bot", [output_not_empty], store=store)

        assert result.all_passed
        assert result.total_traces == 2
        assert result.total_assertions == 1
        assert result.passed == 2
        assert result.failed == 0

    def test_some_fail(self):
        traces = _make_traces(6)
        store = _MockStore(traces)

        result = assert_traces("bot", [output_not_empty], store=store)

        assert not result.all_passed
        # Traces 0 and 3 have empty output
        assert result.failed == 2
        assert result.passed == 4

    def test_multiple_assertions(self):
        traces = _make_traces(5)
        store = _MockStore(traces)

        result = assert_traces(
            "bot",
            [output_not_empty, cost_under_threshold, no_error_spans],
            store=store,
        )

        assert result.total_assertions == 3
        assert result.total_traces == 5
        # 15 total checks (5 traces × 3 assertions)
        assert result.passed + result.failed + result.errors == 15

    def test_exception_counted_as_error(self):
        traces = [{"trace_id": "tr-0", "project": "bot", "spans": []}]
        store = _MockStore(traces)

        result = assert_traces("bot", [always_raises], store=store)

        assert result.errors == 1
        assert result.passed == 0
        assert result.failed == 0
        assert result.results[0].error == "intentional test error"

    def test_last_n_limits_traces(self):
        traces = _make_traces(10)
        store = _MockStore(traces)

        result = assert_traces("bot", [output_not_empty], last=3, store=store)

        assert result.total_traces == 3

    def test_empty_project(self):
        store = _MockStore([])
        result = assert_traces("bot", [output_not_empty], store=store)
        assert result.total_traces == 0
        assert result.all_passed  # no failures = passed

    def test_by_assertion_grouping(self):
        traces = _make_traces(5)
        store = _MockStore(traces)

        result = assert_traces(
            "bot", [output_not_empty, no_error_spans], store=store
        )

        by_name = result.by_assertion()
        assert "output_not_empty" in by_name
        assert "no_error_spans" in by_name
        assert by_name["output_not_empty"]["passed"] + by_name["output_not_empty"]["failed"] == 5

    def test_to_dict(self):
        traces = _make_traces(3)
        store = _MockStore(traces)

        result = assert_traces("bot", [output_not_empty], store=store)
        d = result.to_dict()

        assert d["project"] == "bot"
        assert d["total_traces"] == 3
        assert isinstance(d["results"], list)
        assert all("name" in r and "trace_id" in r for r in d["results"])

    def test_summary_output(self):
        traces = _make_traces(5)
        store = _MockStore(traces)

        result = assert_traces(
            "bot", [output_not_empty, no_error_spans], store=store
        )
        summary = result.summary()

        assert "Assertion Results" in summary
        assert "output_not_empty" in summary
        assert "no_error_spans" in summary
