"""Eval runner — orchestrates running eval cases against an agent."""

from __future__ import annotations

import asyncio
import time
from datetime import datetime, timezone
from pathlib import Path

from hatchdx.agent.config import AgentConfig
from hatchdx.agent.engine import AgentEngine, create_provider, estimate_cost_for_provider
from hatchdx.agent.eval.assertions import EvalContext, evaluate_all_assertions
from hatchdx.agent.eval.mock import MockFailureInterceptor
from hatchdx.agent.eval.models import (
    EvalCase,
    EvalCaseResult,
    EvalRunResult,
    EvalSettings,
    EvalSuite,
)
from hatchdx.agent.providers.base import ModelProvider
from hatchdx.agent.runtime.loop import (
    LoopResult,
    LoopTurn,
    ToolCallEvent,
    _execute_tool_call,
    _parse_timeout,
)
from hatchdx.agent.runtime.servers import ServerManager
from hatchdx.agent.runtime.tools import CollectedTools


# ---------------------------------------------------------------------------
# Mock-aware agent loop
# ---------------------------------------------------------------------------


async def _run_agent_loop_with_mocks(
    provider: ModelProvider,
    server_manager: ServerManager,
    collected_tools: CollectedTools,
    system_prompt: str,
    user_message: str,
    interceptor: MockFailureInterceptor,
    *,
    max_tokens: int = 4096,
    temperature: float = 0.0,
    max_tool_calls: int = 50,
    tool_call_timeout: str = "30s",
    retry_on_error: bool = True,
    max_retries: int = 1,
) -> LoopResult:
    """Run the agent loop with mock failure injection.

    Same as run_agent_loop but intercepts tool calls through the MockFailureInterceptor.
    """
    from hatchdx.agent.runtime.loop import ToolCallLoopError

    timeout_secs = _parse_timeout(tool_call_timeout)

    tool_defs = [
        {
            "name": t.name,
            "description": t.description,
            "inputSchema": t.input_schema,
        }
        for t in collected_tools.tools
    ]

    messages: list[dict] = [{"role": "user", "content": user_message}]
    loop_result = LoopResult()
    tool_call_count = 0

    while True:
        response = await provider.chat(
            messages=messages,
            tools=tool_defs,
            system=system_prompt,
            max_tokens=max_tokens,
            temperature=temperature,
        )

        loop_result.total_input_tokens += response.usage.get("input", 0)
        loop_result.total_output_tokens += response.usage.get("output", 0)

        turn = LoopTurn(response=response)

        if not response.tool_calls:
            loop_result.final_text = response.text
            loop_result.turns.append(turn)
            break

        messages.append({
            "role": "assistant",
            "content": response.text,
            "tool_calls": response.tool_calls,
        })

        for tool_call in response.tool_calls:
            tool_call_count += 1
            if tool_call_count > max_tool_calls:
                loop_result.turns.append(turn)
                loop_result.total_tool_calls = tool_call_count - 1
                raise ToolCallLoopError(
                    f"Exceeded max tool calls ({max_tool_calls})."
                )

            route = collected_tools.get_route(tool_call.name)
            if route is None:
                event = ToolCallEvent(
                    tool_name=tool_call.name,
                    server_name="unknown",
                    arguments=tool_call.arguments,
                    result=f"Unknown tool: '{tool_call.name}'",
                    is_error=True,
                    latency_ms=0,
                )
                turn.tool_events.append(event)
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": event.result,
                    "is_error": True,
                })
                continue

            # Check mock failures BEFORE real execution
            if interceptor.has_failures:
                mock_event = interceptor.maybe_intercept(
                    tool_call.name, route, tool_call.arguments
                )
                if mock_event is not None:
                    turn.tool_events.append(mock_event)
                    messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": mock_event.result,
                        "is_error": True,
                    })
                    continue

            # Real execution
            event = await _execute_tool_call(
                tool_call, route, server_manager, timeout_secs
            )

            if event.is_error and retry_on_error:
                for _ in range(max_retries):
                    retry_event = await _execute_tool_call(
                        tool_call, route, server_manager, timeout_secs
                    )
                    retry_event.was_retry = True
                    if not retry_event.is_error:
                        event = retry_event
                        break
                    event = retry_event

            turn.tool_events.append(event)
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.id,
                "content": event.result,
                "is_error": event.is_error,
            })

        loop_result.turns.append(turn)
        loop_result.total_tool_calls = tool_call_count

    return loop_result


# ---------------------------------------------------------------------------
# Single case runner
# ---------------------------------------------------------------------------


async def run_eval_case(
    case: EvalCase,
    engine: AgentEngine,
    settings: EvalSettings,
) -> EvalCaseResult:
    """Run a single eval case and return the result.

    Args:
        case: The eval case to run.
        engine: A started AgentEngine.
        settings: Eval settings (may override agent defaults).

    Returns:
        EvalCaseResult with assertion results.
    """
    # Resolve settings: eval settings override agent config
    config = engine.config
    temperature = (
        settings.temperature
        if settings.temperature is not None
        else config.model.temperature
    )
    max_tool_calls = (
        settings.max_tool_calls
        if settings.max_tool_calls is not None
        else config.settings.max_tool_calls
    )

    # Build system prompt with optional context injection
    system_prompt = config.system_prompt
    if case.context:
        context_str = "\n".join(f"{k}: {v}" for k, v in case.context.items())
        system_prompt = f"{system_prompt}\n\nContext:\n{context_str}"

    # Set up mock failure interceptor
    interceptor = MockFailureInterceptor(case.mock_failures)

    start_time = time.monotonic()

    try:
        if interceptor.has_failures:
            loop_result = await _run_agent_loop_with_mocks(
                provider=engine.provider,
                server_manager=engine.server_manager,
                collected_tools=engine.collected_tools,
                system_prompt=system_prompt,
                user_message=case.input,
                interceptor=interceptor,
                max_tokens=config.model.max_tokens,
                temperature=temperature,
                max_tool_calls=max_tool_calls,
                tool_call_timeout=config.settings.tool_call_timeout,
                retry_on_error=config.settings.retry_on_error,
                max_retries=config.settings.max_retries,
            )
        else:
            loop_result = await engine.chat(case.input)

    except Exception as e:
        elapsed = (time.monotonic() - start_time) * 1000
        return EvalCaseResult(
            case_name=case.name,
            input=case.input,
            output=None,
            passed=False,
            assertion_results=[],
            tool_calls_json=[],
            latency_ms=elapsed,
            tokens_used=0,
            cost_usd=0.0,
            error_message=str(e),
        )

    elapsed_ms = (time.monotonic() - start_time) * 1000

    # Cost calculation
    cost = estimate_cost_for_provider(
        config.model.provider,
        config.model.model,
        loop_result.total_input_tokens,
        loop_result.total_output_tokens,
    )

    # Build tool calls JSON for storage
    tool_calls_json = [
        {
            "name": e.tool_name,
            "server": e.server_name,
            "arguments": e.arguments,
            "result": e.result[:200],  # Truncate for storage
            "is_error": e.is_error,
            "latency_ms": round(e.latency_ms, 1),
        }
        for e in loop_result.all_tool_events
    ]

    # Evaluate assertions
    ctx = EvalContext(
        loop_result=loop_result,
        latency_ms=elapsed_ms,
        cost_usd=cost,
    )
    assertion_results = evaluate_all_assertions(case.assertions, ctx)

    all_passed = all(ar.passed for ar in assertion_results)
    error_msg = None
    if not all_passed:
        failed = [ar for ar in assertion_results if not ar.passed]
        error_msg = "; ".join(f"[{ar.assertion_type}] {ar.message}" for ar in failed)

    return EvalCaseResult(
        case_name=case.name,
        input=case.input,
        output=loop_result.final_text,
        passed=all_passed,
        assertion_results=assertion_results,
        tool_calls_json=tool_calls_json,
        latency_ms=elapsed_ms,
        tokens_used=loop_result.total_input_tokens + loop_result.total_output_tokens,
        cost_usd=cost,
        error_message=error_msg,
    )


# ---------------------------------------------------------------------------
# Suite runner
# ---------------------------------------------------------------------------


async def run_eval_suite(
    suite: EvalSuite,
    engine: AgentEngine,
    *,
    tag_filter: str | None = None,
    on_case_complete: object | None = None,
) -> EvalRunResult:
    """Run all cases in an eval suite sequentially.

    Args:
        suite: The eval suite to run.
        engine: A started AgentEngine.
        tag_filter: If set, only run cases with this tag.
        on_case_complete: Optional async callback(EvalCaseResult) for progress.

    Returns:
        EvalRunResult with all case results.
    """
    run_result = EvalRunResult(
        suite_name=suite.name,
        agent_name=suite.agent,
        model=suite.settings.model or engine.config.model.model,
        temperature=(
            suite.settings.temperature
            if suite.settings.temperature is not None
            else engine.config.model.temperature
        ),
        started_at=datetime.now(timezone.utc),
    )

    cases = suite.cases
    if tag_filter:
        cases = [c for c in cases if tag_filter in c.tags]

    for case in cases:
        case_result = await run_eval_case(case, engine, suite.settings)
        run_result.case_results.append(case_result)

        if on_case_complete and asyncio.iscoroutinefunction(on_case_complete):
            await on_case_complete(case_result)

    run_result.completed_at = datetime.now(timezone.utc)
    return run_result


async def run_eval_suite_multi(
    suite: EvalSuite,
    engine: AgentEngine,
    *,
    runs: int = 1,
    tag_filter: str | None = None,
    on_case_complete: object | None = None,
) -> list[EvalRunResult]:
    """Run an eval suite multiple times and return all results.

    Args:
        suite: The eval suite to run.
        engine: A started AgentEngine.
        runs: Number of times to run the suite.
        tag_filter: If set, only run cases with this tag.
        on_case_complete: Optional async callback(EvalCaseResult) for progress.

    Returns:
        List of EvalRunResult, one per run.
    """
    results = []
    for _ in range(runs):
        result = await run_eval_suite(
            suite, engine, tag_filter=tag_filter, on_case_complete=on_case_complete
        )
        results.append(result)
    return results


# ---------------------------------------------------------------------------
# Cross-model eval
# ---------------------------------------------------------------------------


async def run_eval_cross_model(
    suite: EvalSuite,
    engine: AgentEngine,
    models: list[str],
    *,
    tag_filter: str | None = None,
    on_case_complete: object | None = None,
    on_model_start: object | None = None,
) -> list[EvalRunResult]:
    """Run an eval suite against multiple models for comparison.

    For each model, temporarily swaps the engine's provider and re-runs the suite.
    The engine must already be started (servers running).

    Args:
        suite: The eval suite to run.
        engine: A started AgentEngine.
        models: List of model identifiers to test (e.g., ["claude-sonnet-4-20250514", "gpt-4o"]).
        tag_filter: If set, only run cases with this tag.
        on_case_complete: Optional async callback(EvalCaseResult).
        on_model_start: Optional async callback(model_name: str).

    Returns:
        List of EvalRunResult, one per model.
    """
    from hatchdx.agent.config import ModelConfig

    results = []
    original_provider = engine._provider
    original_model_config = engine.config.model

    for model_spec in models:
        # Parse model spec: "provider:model" or just "model" (keeps current provider)
        if ":" in model_spec and not model_spec.startswith("http"):
            provider_name, model_name = model_spec.split(":", 1)
        else:
            provider_name = original_model_config.provider
            model_name = model_spec

        if on_model_start and asyncio.iscoroutinefunction(on_model_start):
            await on_model_start(f"{provider_name}:{model_name}")

        # Create a temporary config with the new model
        temp_model_config = ModelConfig(
            provider=provider_name,
            model=model_name,
            max_tokens=original_model_config.max_tokens,
            temperature=original_model_config.temperature,
            base_url=original_model_config.base_url,
        )

        # Swap provider temporarily
        from hatchdx.agent.config import AgentConfig

        temp_config = AgentConfig(
            name=engine.config.name,
            description=engine.config.description,
            model=temp_model_config,
            system_prompt=engine.config.system_prompt,
            servers=engine.config.servers,
            version=engine.config.version,
            system_prompt_file=engine.config.system_prompt_file,
            tool_filter=engine.config.tool_filter,
            settings=engine.config.settings,
        )

        try:
            temp_provider = create_provider(temp_config)
        except Exception as e:
            # Create a failed result for this model
            run_result = EvalRunResult(
                suite_name=suite.name,
                agent_name=suite.agent,
                model=f"{provider_name}:{model_name}",
                temperature=original_model_config.temperature,
                started_at=datetime.now(timezone.utc),
                completed_at=datetime.now(timezone.utc),
            )
            results.append(run_result)
            continue

        # Temporarily swap engine internals
        engine._provider = temp_provider
        engine.config = temp_config

        try:
            run_result = await run_eval_suite(
                suite, engine, tag_filter=tag_filter, on_case_complete=on_case_complete
            )
            # Override the model field to include provider for clarity
            run_result.model = f"{provider_name}:{model_name}"
            results.append(run_result)
        except Exception:
            # Record the error but continue with other models
            run_result = EvalRunResult(
                suite_name=suite.name,
                agent_name=suite.agent,
                model=f"{provider_name}:{model_name}",
                temperature=original_model_config.temperature,
                started_at=datetime.now(timezone.utc),
                completed_at=datetime.now(timezone.utc),
            )
            results.append(run_result)

    # Restore original state
    engine._provider = original_provider
    engine.config = AgentConfig(
        name=engine.config.name,
        description=engine.config.description,
        model=original_model_config,
        system_prompt=engine.config.system_prompt,
        servers=engine.config.servers,
        version=engine.config.version,
        system_prompt_file=engine.config.system_prompt_file,
        tool_filter=engine.config.tool_filter,
        settings=engine.config.settings,
    )

    return results
