"""Assertion runner for Bandito traces.

Assertions are Python functions: trace_dict → bool.
Truthy = pass, falsy = fail, exception = error.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Callable


@dataclass
class AssertionResult:
    """Result of a single assertion on a single trace."""

    name: str
    trace_id: str
    passed: bool
    error: str | None = None


@dataclass
class AssertionRunResult:
    """Aggregate result of running assertions against traces."""

    project: str
    total_traces: int = 0
    total_assertions: int = 0
    passed: int = 0
    failed: int = 0
    errors: int = 0
    results: list[AssertionResult] = field(default_factory=list)

    @property
    def all_passed(self) -> bool:
        return self.failed == 0 and self.errors == 0

    def by_assertion(self) -> dict[str, dict[str, int]]:
        """Group results by assertion name → {passed, failed, errors}."""
        groups: dict[str, dict[str, int]] = {}
        for r in self.results:
            if r.name not in groups:
                groups[r.name] = {"passed": 0, "failed": 0, "errors": 0}
            if r.error:
                groups[r.name]["errors"] += 1
            elif r.passed:
                groups[r.name]["passed"] += 1
            else:
                groups[r.name]["failed"] += 1
        return groups

    def to_dict(self) -> dict[str, Any]:
        """Serialize for backend storage."""
        return {
            "project": self.project,
            "total_traces": self.total_traces,
            "total_assertions": self.total_assertions,
            "passed": self.passed,
            "failed": self.failed,
            "errors": self.errors,
            "results": [
                {
                    "name": r.name,
                    "trace_id": r.trace_id,
                    "passed": r.passed,
                    "error": r.error,
                }
                for r in self.results
            ],
        }

    def summary(self) -> str:
        """Terminal-formatted summary."""
        lines = [
            "",
            "  Assertion Results",
            f"  {'─' * 50}",
            "",
            f"  Traces:     {self.total_traces}",
            f"  Assertions: {self.total_assertions}",
            "",
        ]

        by_name = self.by_assertion()
        for name, counts in sorted(by_name.items()):
            total = counts["passed"] + counts["failed"] + counts["errors"]
            if counts["failed"] == 0 and counts["errors"] == 0:
                mark = "✓"
            else:
                mark = "✗"
            parts = [f"{counts['passed']}/{total} passed"]
            if counts["failed"]:
                parts.append(f"{counts['failed']} failed")
            if counts["errors"]:
                parts.append(f"{counts['errors']} errors")
            lines.append(f"  {mark} {name:<30} {' · '.join(parts)}")

        lines.append("")
        if self.all_passed:
            lines.append(f"  All assertions passed.")
        else:
            lines.append(
                f"  {self.failed} failed, {self.errors} errors "
                f"out of {self.passed + self.failed + self.errors} checks."
            )

            # Show first few failures
            failures = [r for r in self.results if not r.passed and not r.error]
            if failures:
                lines.append("")
                lines.append("  Failures:")
                for r in failures[:5]:
                    lines.append(f"    {r.name} · {r.trace_id}")
                if len(failures) > 5:
                    lines.append(f"    ... and {len(failures) - 5} more")

            errs = [r for r in self.results if r.error]
            if errs:
                lines.append("")
                lines.append("  Errors:")
                for r in errs[:3]:
                    lines.append(f"    {r.name} · {r.trace_id}: {r.error}")
                if len(errs) > 3:
                    lines.append(f"    ... and {len(errs) - 3} more")

        lines.append("")
        return "\n".join(lines)


def assert_traces(
    project: str,
    assertions: list[Callable[[dict[str, Any]], bool]],
    *,
    last: int | None = None,
    store: Any = None,
) -> AssertionRunResult:
    """Run assertions against traces.

    Each assertion is a function: trace_dict → bool.
    Truthy = pass, falsy = fail, exception = error.

    Args:
        project: Project name.
        assertions: List of assertion functions.
        last: Only check the last N traces.
        store: StoreBase instance (resolved from config if None).

    Returns:
        AssertionRunResult with per-assertion, per-trace results.
    """
    if store is None:
        from .storage import get_store

        store = get_store(project=project)

    traces = store.list_traces_filtered(project=project, limit=last)

    result = AssertionRunResult(
        project=project,
        total_traces=len(traces),
        total_assertions=len(assertions),
    )

    for trace in traces:
        trace_id = trace.get("trace_id", "?")
        for fn in assertions:
            name = fn.__name__
            try:
                passed = bool(fn(trace))
            except Exception as e:
                result.errors += 1
                result.results.append(
                    AssertionResult(
                        name=name,
                        trace_id=trace_id,
                        passed=False,
                        error=str(e),
                    )
                )
                continue

            if passed:
                result.passed += 1
            else:
                result.failed += 1

            result.results.append(
                AssertionResult(
                    name=name,
                    trace_id=trace_id,
                    passed=passed,
                )
            )

    return result
