"""Eval data models — dataclasses for suites, cases, results, and assertions."""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path

import yaml

from hatchdx import HdxError


# ---------------------------------------------------------------------------
# Errors
# ---------------------------------------------------------------------------


class EvalConfigError(HdxError):
    """Invalid eval suite configuration."""


# ---------------------------------------------------------------------------
# Eval suite models
# ---------------------------------------------------------------------------


@dataclass
class EvalSettings:
    """Optional settings overrides for an eval suite."""

    model: str | None = None
    temperature: float | None = None
    max_tool_calls: int | None = None


@dataclass
class MockFailure:
    """Mock failure injection for an eval case."""

    server: str
    error: str
    tools: list[str] = field(default_factory=list)
    fail_count: int = 0  # 0 = always fail


@dataclass
class AssertionConfig:
    """A single assertion to check against agent output."""

    type: str
    # Tool behavior
    tools: list[str] = field(default_factory=list)
    sequence: list[str] = field(default_factory=list)
    count: int = 0
    # Output content
    values: list[str] = field(default_factory=list)
    pattern: str = ""
    format: str = ""
    # Performance
    seconds: float = 0.0
    tokens: int = 0
    amount: float = 0.0


@dataclass
class EvalCase:
    """A single eval case within a suite."""

    name: str
    input: str
    assertions: list[AssertionConfig]
    context: dict = field(default_factory=dict)
    tags: list[str] = field(default_factory=list)
    timeout: str | None = None
    mock_failures: list[MockFailure] = field(default_factory=list)


@dataclass
class EvalSuite:
    """A complete eval suite parsed from YAML."""

    name: str
    description: str
    agent: str
    cases: list[EvalCase]
    settings: EvalSettings = field(default_factory=EvalSettings)


# ---------------------------------------------------------------------------
# Result models
# ---------------------------------------------------------------------------


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

    assertion_type: str
    passed: bool
    message: str
    details: dict = field(default_factory=dict)


@dataclass
class EvalCaseResult:
    """Result of running a single eval case."""

    case_name: str
    input: str
    output: str | None
    passed: bool
    assertion_results: list[AssertionResult]
    tool_calls_json: list[dict]
    latency_ms: float
    tokens_used: int
    cost_usd: float
    error_message: str | None = None


@dataclass
class EvalRunResult:
    """Result of running a complete eval suite."""

    suite_name: str
    agent_name: str
    model: str
    temperature: float
    started_at: datetime
    completed_at: datetime | None = None
    case_results: list[EvalCaseResult] = field(default_factory=list)

    @property
    def total_cases(self) -> int:
        return len(self.case_results)

    @property
    def passed(self) -> int:
        return sum(1 for c in self.case_results if c.passed)

    @property
    def failed(self) -> int:
        return sum(1 for c in self.case_results if not c.passed)

    @property
    def total_cost(self) -> float:
        return sum(c.cost_usd for c in self.case_results)

    @property
    def total_tokens(self) -> int:
        return sum(c.tokens_used for c in self.case_results)


# ---------------------------------------------------------------------------
# Assertion type registry
# ---------------------------------------------------------------------------

VALID_ASSERTION_TYPES = {
    # Tool behavior
    "tool_called",
    "tool_sequence_includes",
    "tool_not_called",
    "no_repeated_tool_calls",
    "min_tool_calls",
    "max_tool_calls",
    # Output content
    "output_contains",
    "output_not_contains",
    "output_matches",
    "output_format",
    # Data integrity
    "cites_data_source",
    "no_hallucinated_stats",
    # Performance
    "max_latency_seconds",
    "max_tokens_used",
    "max_cost_usd",
}


# ---------------------------------------------------------------------------
# Parsing
# ---------------------------------------------------------------------------


def _parse_assertion(data: dict, case_name: str, index: int) -> AssertionConfig:
    """Parse a single assertion from YAML data."""
    atype = data.get("type")
    if not atype:
        raise EvalConfigError(
            f"Case '{case_name}', assertion {index}: 'type' is required"
        )
    if atype not in VALID_ASSERTION_TYPES:
        raise EvalConfigError(
            f"Case '{case_name}', assertion {index}: unknown type '{atype}'\n"
            f"Valid types: {', '.join(sorted(VALID_ASSERTION_TYPES))}"
        )

    return AssertionConfig(
        type=atype,
        tools=data.get("tools", []),
        sequence=data.get("sequence", []),
        count=data.get("count", 0),
        values=data.get("values", []),
        pattern=data.get("pattern", ""),
        format=data.get("format", ""),
        seconds=data.get("seconds", 0.0),
        tokens=data.get("tokens", 0),
        amount=data.get("amount", 0.0),
    )


def _parse_mock_failure(data: dict, case_name: str, index: int) -> MockFailure:
    """Parse a single mock failure from YAML data."""
    server = data.get("server")
    if not server:
        raise EvalConfigError(
            f"Case '{case_name}', mock_failure {index}: 'server' is required"
        )
    error = data.get("error")
    if not error:
        raise EvalConfigError(
            f"Case '{case_name}', mock_failure {index}: 'error' is required"
        )
    return MockFailure(
        server=server,
        error=error,
        tools=data.get("tools", []),
        fail_count=data.get("fail_count", 0),
    )


def _parse_case(data: dict, index: int) -> EvalCase:
    """Parse a single eval case from YAML data."""
    name = data.get("name")
    if not name:
        raise EvalConfigError(f"Case at index {index}: 'name' is required")

    input_text = data.get("input")
    if not input_text:
        raise EvalConfigError(f"Case '{name}': 'input' is required")

    assertions_data = data.get("assertions", [])
    if not assertions_data:
        raise EvalConfigError(f"Case '{name}': at least one assertion is required")

    assertions = [
        _parse_assertion(a, name, i)
        for i, a in enumerate(assertions_data)
    ]

    mock_failures = [
        _parse_mock_failure(m, name, i)
        for i, m in enumerate(data.get("mock_failures", []))
    ]

    return EvalCase(
        name=name,
        input=input_text,
        assertions=assertions,
        context=data.get("context", {}),
        tags=data.get("tags", []),
        timeout=data.get("timeout"),
        mock_failures=mock_failures,
    )


def _parse_settings(data: dict | None) -> EvalSettings:
    """Parse optional eval settings overrides."""
    if not data:
        return EvalSettings()
    return EvalSettings(
        model=data.get("model"),
        temperature=data.get("temperature"),
        max_tool_calls=data.get("max_tool_calls"),
    )


def load_eval_suite(path: Path) -> EvalSuite:
    """Load and validate an eval suite from a YAML file.

    Args:
        path: Path to the eval YAML file.

    Returns:
        Validated EvalSuite.

    Raises:
        EvalConfigError: If the suite is invalid.
        FileNotFoundError: If the file doesn't exist.
    """
    if not path.exists():
        raise FileNotFoundError(f"Eval suite not found: {path}")

    try:
        raw = yaml.safe_load(path.read_text())
    except yaml.YAMLError as e:
        raise EvalConfigError(f"Invalid YAML in {path}: {e}") from e

    if not isinstance(raw, dict):
        raise EvalConfigError(f"Expected a YAML mapping in {path}")

    name = raw.get("name")
    if not name:
        raise EvalConfigError("Eval suite 'name' is required")

    description = raw.get("description", "")

    agent = raw.get("agent")
    if not agent:
        raise EvalConfigError("Eval suite 'agent' is required")

    cases_data = raw.get("cases", [])
    if not cases_data:
        raise EvalConfigError("Eval suite must have at least one case")
    if not isinstance(cases_data, list):
        raise EvalConfigError("'cases' must be a list")

    cases = [_parse_case(c, i) for i, c in enumerate(cases_data)]

    settings = _parse_settings(raw.get("settings"))

    return EvalSuite(
        name=name,
        description=description,
        agent=agent,
        cases=cases,
        settings=settings,
    )
