"""Assertion engine — evaluates assertions against agent loop results."""

from __future__ import annotations

import json
import re

from hatchdx.agent.eval.models import AssertionConfig, AssertionResult
from hatchdx.agent.runtime.loop import LoopResult, ToolCallEvent


# ---------------------------------------------------------------------------
# Context object passed to each assertion checker
# ---------------------------------------------------------------------------


class EvalContext:
    """Holds all data needed to evaluate assertions for a single eval case."""

    def __init__(
        self,
        loop_result: LoopResult,
        latency_ms: float,
        cost_usd: float,
    ) -> None:
        self.loop_result = loop_result
        self.output = loop_result.final_text or ""
        self.tool_events = loop_result.all_tool_events
        self.latency_ms = latency_ms
        self.cost_usd = cost_usd
        self.total_tokens = (
            loop_result.total_input_tokens + loop_result.total_output_tokens
        )

    @property
    def tool_names_called(self) -> list[str]:
        return [e.tool_name for e in self.tool_events]

    @property
    def tool_results(self) -> list[str]:
        return [e.result for e in self.tool_events if not e.is_error]


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

_ASSERTION_CHECKERS: dict[str, callable] = {}


def _register(name: str):
    """Decorator to register an assertion checker."""
    def decorator(fn):
        _ASSERTION_CHECKERS[name] = fn
        return fn
    return decorator


def evaluate_assertion(
    assertion: AssertionConfig,
    ctx: EvalContext,
) -> AssertionResult:
    """Evaluate a single assertion against an EvalContext.

    Returns an AssertionResult with pass/fail and a descriptive message.
    """
    checker = _ASSERTION_CHECKERS.get(assertion.type)
    if checker is None:
        return AssertionResult(
            assertion_type=assertion.type,
            passed=False,
            message=f"Unknown assertion type: {assertion.type}",
        )
    return checker(assertion, ctx)


def evaluate_all_assertions(
    assertions: list[AssertionConfig],
    ctx: EvalContext,
) -> list[AssertionResult]:
    """Evaluate all assertions for a case. Returns list of results."""
    return [evaluate_assertion(a, ctx) for a in assertions]


# ---------------------------------------------------------------------------
# Tool behavior assertions
# ---------------------------------------------------------------------------


@_register("tool_called")
def _check_tool_called(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """All listed tools must have been called (order doesn't matter)."""
    called = set(ctx.tool_names_called)
    missing = [t for t in a.tools if t not in called]
    if missing:
        return AssertionResult(
            assertion_type="tool_called",
            passed=False,
            message=f"Tools not called: {', '.join(missing)}",
            details={"missing": missing, "called": list(called)},
        )
    return AssertionResult(
        assertion_type="tool_called",
        passed=True,
        message=f"All required tools called: {', '.join(a.tools)}",
    )


@_register("tool_sequence_includes")
def _check_tool_sequence(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """Tools must appear in the specified order (other tools can be interspersed)."""
    called = ctx.tool_names_called
    seq = a.sequence
    seq_idx = 0

    for tool_name in called:
        if seq_idx < len(seq) and tool_name == seq[seq_idx]:
            seq_idx += 1

    if seq_idx == len(seq):
        return AssertionResult(
            assertion_type="tool_sequence_includes",
            passed=True,
            message=f"Sequence found: {' → '.join(seq)}",
        )

    return AssertionResult(
        assertion_type="tool_sequence_includes",
        passed=False,
        message=f"Sequence not found. Expected: {' → '.join(seq)}. "
                f"Got: {' → '.join(called) or '(no tools called)'}",
        details={"expected": seq, "actual": called, "matched_up_to": seq_idx},
    )


@_register("tool_not_called")
def _check_tool_not_called(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """None of the listed tools should have been called."""
    called = set(ctx.tool_names_called)
    found = [t for t in a.tools if t in called]
    if found:
        return AssertionResult(
            assertion_type="tool_not_called",
            passed=False,
            message=f"Prohibited tools were called: {', '.join(found)}",
            details={"prohibited_called": found},
        )
    return AssertionResult(
        assertion_type="tool_not_called",
        passed=True,
        message="No prohibited tools called",
    )


@_register("no_repeated_tool_calls")
def _check_no_repeated(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """Fails if the exact same tool+arguments combination is called twice."""
    seen: set[str] = set()
    for event in ctx.tool_events:
        key = f"{event.tool_name}:{json.dumps(event.arguments, sort_keys=True)}"
        if key in seen:
            return AssertionResult(
                assertion_type="no_repeated_tool_calls",
                passed=False,
                message=f"Repeated tool call: {event.tool_name} with same arguments",
                details={"repeated_tool": event.tool_name, "arguments": event.arguments},
            )
        seen.add(key)
    return AssertionResult(
        assertion_type="no_repeated_tool_calls",
        passed=True,
        message="No repeated tool calls detected",
    )


@_register("min_tool_calls")
def _check_min_tool_calls(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """Agent must use at least N tools."""
    actual = len(ctx.tool_events)
    if actual < a.count:
        return AssertionResult(
            assertion_type="min_tool_calls",
            passed=False,
            message=f"Too few tool calls: {actual} < {a.count}",
            details={"actual": actual, "minimum": a.count},
        )
    return AssertionResult(
        assertion_type="min_tool_calls",
        passed=True,
        message=f"Tool call count OK: {actual} >= {a.count}",
    )


@_register("max_tool_calls")
def _check_max_tool_calls(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """Agent must not use more than N tools."""
    actual = len(ctx.tool_events)
    if actual > a.count:
        return AssertionResult(
            assertion_type="max_tool_calls",
            passed=False,
            message=f"Too many tool calls: {actual} > {a.count}",
            details={"actual": actual, "maximum": a.count},
        )
    return AssertionResult(
        assertion_type="max_tool_calls",
        passed=True,
        message=f"Tool call count OK: {actual} <= {a.count}",
    )


# ---------------------------------------------------------------------------
# Output content assertions
# ---------------------------------------------------------------------------


@_register("output_contains")
def _check_output_contains(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """All values must appear in the final text response (case-insensitive)."""
    output_lower = ctx.output.lower()
    missing = [v for v in a.values if v.lower() not in output_lower]
    if missing:
        return AssertionResult(
            assertion_type="output_contains",
            passed=False,
            message=f"Output missing: {', '.join(repr(v) for v in missing)}",
            details={"missing": missing},
        )
    return AssertionResult(
        assertion_type="output_contains",
        passed=True,
        message=f"Output contains all required values",
    )


@_register("output_not_contains")
def _check_output_not_contains(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """None of the values should appear in the output (case-insensitive)."""
    output_lower = ctx.output.lower()
    found = [v for v in a.values if v.lower() in output_lower]
    if found:
        return AssertionResult(
            assertion_type="output_not_contains",
            passed=False,
            message=f"Output contains prohibited values: {', '.join(repr(v) for v in found)}",
            details={"found": found},
        )
    return AssertionResult(
        assertion_type="output_not_contains",
        passed=True,
        message="Output does not contain prohibited values",
    )


@_register("output_matches")
def _check_output_matches(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """Response must match the given regex pattern."""
    try:
        match = re.search(a.pattern, ctx.output)
    except re.error as e:
        return AssertionResult(
            assertion_type="output_matches",
            passed=False,
            message=f"Invalid regex pattern: {e}",
        )

    if match:
        return AssertionResult(
            assertion_type="output_matches",
            passed=True,
            message=f"Output matches pattern: {a.pattern}",
            details={"matched": match.group()},
        )
    return AssertionResult(
        assertion_type="output_matches",
        passed=False,
        message=f"Output does not match pattern: {a.pattern}",
    )


@_register("output_format")
def _check_output_format(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """Check structural format of the output."""
    fmt = a.format
    output = ctx.output

    if fmt == "structured_recommendation":
        # Must have headers (## or **bold**) and bullet points
        has_headers = bool(re.search(r"(^##\s|^\*\*[^*]+\*\*)", output, re.MULTILINE))
        has_bullets = bool(re.search(r"^[\s]*[-*]\s", output, re.MULTILINE))
        if has_headers and has_bullets:
            return AssertionResult(
                assertion_type="output_format",
                passed=True,
                message="Output has structured recommendation format",
            )
        missing = []
        if not has_headers:
            missing.append("headers")
        if not has_bullets:
            missing.append("bullet points")
        return AssertionResult(
            assertion_type="output_format",
            passed=False,
            message=f"Output missing: {', '.join(missing)}",
        )

    elif fmt == "data_table":
        # Must contain pipe-separated table or aligned columns
        has_table = bool(re.search(r"\|.*\|.*\|", output))
        if has_table:
            return AssertionResult(
                assertion_type="output_format",
                passed=True,
                message="Output contains data table",
            )
        return AssertionResult(
            assertion_type="output_format",
            passed=False,
            message="Output does not contain a data table",
        )

    elif fmt == "summary":
        # Concise paragraphs, no tool call artifacts
        has_artifacts = bool(re.search(r"(tool_use|tool_result|<tool|```json\s*\{)", output))
        if has_artifacts:
            return AssertionResult(
                assertion_type="output_format",
                passed=False,
                message="Output contains tool call artifacts",
            )
        return AssertionResult(
            assertion_type="output_format",
            passed=True,
            message="Output is a clean summary",
        )

    return AssertionResult(
        assertion_type="output_format",
        passed=False,
        message=f"Unknown format: '{fmt}'",
    )


# ---------------------------------------------------------------------------
# Data integrity assertions
# ---------------------------------------------------------------------------


@_register("cites_data_source")
def _check_cites_data(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """Checks that the agent's response references data from tool call results."""
    if not ctx.tool_results:
        return AssertionResult(
            assertion_type="cites_data_source",
            passed=False,
            message="No tool results to cross-reference",
        )

    # Extract significant words/numbers from tool results
    tool_data_tokens: set[str] = set()
    for result in ctx.tool_results:
        # Extract numbers and multi-char words from tool results
        tokens = re.findall(r"\b\d+\.?\d*\b|\b[a-zA-Z]{3,}\b", result)
        tool_data_tokens.update(t.lower() for t in tokens)

    # Check if the output references any of the tool data
    output_tokens = set(
        t.lower() for t in re.findall(r"\b\d+\.?\d*\b|\b[a-zA-Z]{3,}\b", ctx.output)
    )

    overlap = tool_data_tokens & output_tokens
    if overlap:
        return AssertionResult(
            assertion_type="cites_data_source",
            passed=True,
            message=f"Output references tool data ({len(overlap)} shared tokens)",
            details={"overlap_count": len(overlap)},
        )

    return AssertionResult(
        assertion_type="cites_data_source",
        passed=False,
        message="Output does not reference any data from tool results",
    )


@_register("no_hallucinated_stats")
def _check_no_hallucinated_stats(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """Cross-references numbers in output with numbers from tool results."""
    if not ctx.tool_results:
        # No tool results to check against — vacuously true
        return AssertionResult(
            assertion_type="no_hallucinated_stats",
            passed=True,
            message="No tool results to cross-reference (vacuously true)",
        )

    # Extract numbers from tool results
    tool_numbers: set[str] = set()
    for result in ctx.tool_results:
        numbers = re.findall(r"\b\d+\.?\d*\b", result)
        tool_numbers.update(numbers)

    # Extract numbers from output
    output_numbers = re.findall(r"\b\d+\.?\d*\b", ctx.output)

    # Filter out trivially common numbers (1, 2, etc.)
    significant_output_numbers = [
        n for n in output_numbers
        if float(n) > 10 or "." in n
    ]

    hallucinated = [n for n in significant_output_numbers if n not in tool_numbers]

    if hallucinated:
        return AssertionResult(
            assertion_type="no_hallucinated_stats",
            passed=False,
            message=f"Potentially hallucinated numbers: {', '.join(hallucinated[:5])}",
            details={"hallucinated": hallucinated, "tool_numbers": list(tool_numbers)},
        )

    return AssertionResult(
        assertion_type="no_hallucinated_stats",
        passed=True,
        message="All significant numbers found in tool results",
    )


# ---------------------------------------------------------------------------
# Performance assertions
# ---------------------------------------------------------------------------


@_register("max_latency_seconds")
def _check_max_latency(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """Total wall-clock time from input to final response."""
    actual_seconds = ctx.latency_ms / 1000.0
    if actual_seconds > a.seconds:
        return AssertionResult(
            assertion_type="max_latency_seconds",
            passed=False,
            message=f"Latency {actual_seconds:.1f}s > {a.seconds}s limit",
            details={"actual_seconds": actual_seconds, "limit": a.seconds},
        )
    return AssertionResult(
        assertion_type="max_latency_seconds",
        passed=True,
        message=f"Latency OK: {actual_seconds:.1f}s <= {a.seconds}s",
    )


@_register("max_tokens_used")
def _check_max_tokens(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """Total tokens (input + output) across all model calls."""
    if ctx.total_tokens > a.tokens:
        return AssertionResult(
            assertion_type="max_tokens_used",
            passed=False,
            message=f"Token usage {ctx.total_tokens:,} > {a.tokens:,} limit",
            details={"actual": ctx.total_tokens, "limit": a.tokens},
        )
    return AssertionResult(
        assertion_type="max_tokens_used",
        passed=True,
        message=f"Token usage OK: {ctx.total_tokens:,} <= {a.tokens:,}",
    )


@_register("max_cost_usd")
def _check_max_cost(a: AssertionConfig, ctx: EvalContext) -> AssertionResult:
    """Total estimated cost for this eval case."""
    if ctx.cost_usd > a.amount:
        return AssertionResult(
            assertion_type="max_cost_usd",
            passed=False,
            message=f"Cost ${ctx.cost_usd:.4f} > ${a.amount:.4f} limit",
            details={"actual": ctx.cost_usd, "limit": a.amount},
        )
    return AssertionResult(
        assertion_type="max_cost_usd",
        passed=True,
        message=f"Cost OK: ${ctx.cost_usd:.4f} <= ${a.amount:.4f}",
    )
