"""Capture mode mixin for EvalSuite.

This module provides the capture functionality as a mixin class,
keeping it separate from the main evaluation logic in eval.py.
"""

from __future__ import annotations

import asyncio
import random
from typing import TYPE_CHECKING, Any

from arcade_evals._evalsuite._types import _resolve_seed_spec
from arcade_evals.capture import CapturedCase, CapturedRun, CapturedToolCall, CaptureResult

if TYPE_CHECKING:
    from arcade_evals._evalsuite._comparative import ComparativeCaseBuilder
    from arcade_evals._evalsuite._providers import ProviderName
    from arcade_evals._evalsuite._tool_registry import EvalSuiteToolRegistry
    from arcade_evals._evalsuite._tracks import TrackManager
    from arcade_evals._evalsuite._types import EvalRubric
    from arcade_evals.eval import EvalCase


class _EvalSuiteCaptureMixin:
    """Mixin providing capture mode functionality for EvalSuite."""

    # These attributes are defined in EvalSuite
    name: str
    cases: list[EvalCase]
    max_concurrent: int
    rubric: EvalRubric
    _internal_registry: EvalSuiteToolRegistry | None
    _comparative_case_builders: list[ComparativeCaseBuilder]
    _track_manager: TrackManager

    # These methods are defined in EvalSuite
    async def _run_openai(
        self,
        client: Any,
        model: str,
        case: EvalCase,
        registry: EvalSuiteToolRegistry | None = None,
        seed: int | None = None,
    ) -> list[tuple[str, dict[str, Any]]]:
        raise NotImplementedError  # Implemented in EvalSuite

    async def _run_anthropic(
        self,
        client: Any,
        model: str,
        case: EvalCase,
        registry: EvalSuiteToolRegistry | None = None,
    ) -> list[tuple[str, dict[str, Any]]]:
        raise NotImplementedError  # Implemented in EvalSuite

    def _process_tool_calls(
        self,
        tool_calls: list[tuple[str, dict[str, Any]]],
        registry: EvalSuiteToolRegistry | None = None,
    ) -> list[tuple[str, dict[str, Any]]]:
        raise NotImplementedError  # Implemented in EvalSuite

    def _create_eval_case(self, *args: Any, **kwargs: Any) -> EvalCase:
        raise NotImplementedError  # Implemented in EvalSuite

    async def capture(
        self,
        client: Any,  # AsyncOpenAI | AsyncAnthropic
        model: str,
        provider: ProviderName = "openai",
        include_context: bool = False,
        num_runs: int = 1,
        seed: str | int | None = "constant",
    ) -> CaptureResult:
        """
        Run the evaluation suite in capture mode - records tool calls without scoring.

        Capture mode runs each case and records the tool calls made by the model,
        without evaluating or scoring them. This is useful for:
        - Generating expected tool calls for new test cases
        - Debugging model behavior
        - Creating baseline recordings

        Handles both regular cases and comparative cases. For comparative cases,
        each track is captured separately with its own tool registry.

        Args:
            client: The LLM client instance (AsyncOpenAI or AsyncAnthropic).
            model: The model to use.
            provider: The provider name ("openai" or "anthropic").
            include_context: Whether to include system_message and additional_messages
                           in the output.
            num_runs: Number of runs per case.
            seed: Seed policy ("constant", "random", or an integer seed).

        Returns:
            A CaptureResult containing all captured tool calls.
        """
        if num_runs < 1:
            raise ValueError("num_runs must be >= 1")

        all_captured: list[CapturedCase] = []
        semaphore = asyncio.Semaphore(self.max_concurrent)

        async def capture_case(
            case: EvalCase,
            registry: EvalSuiteToolRegistry | None = None,
            track: str | None = None,
        ) -> CapturedCase:
            """Capture a case using the specified registry."""
            async with semaphore:
                use_registry = registry or self._internal_registry
                if use_registry is None or use_registry.tool_count() == 0:
                    raise ValueError(
                        "No tools registered. Use add_* convenience methods or pass catalog=ToolCatalog."
                    )

                seed_policy, seed_value = _resolve_seed_spec(seed)
                if provider == "openai":
                    if seed_policy == "random":
                        run_seeds: list[int | None] = [
                            random.randint(0, 2**31 - 1)  # noqa: S311
                            for _ in range(num_runs)
                        ]
                    else:
                        run_seeds = [seed_value for _ in range(num_runs)]
                else:
                    run_seeds = [None for _ in range(num_runs)]

                runs: list[CapturedRun] = []
                for run_index in range(num_runs):
                    run_seed = run_seeds[run_index]
                    # Get tool calls based on provider
                    if provider == "anthropic":
                        predicted_args = await self._run_anthropic(
                            client, model, case, registry=use_registry
                        )
                    else:
                        predicted_args = await self._run_openai(
                            client, model, case, registry=use_registry, seed=run_seed
                        )

                    # Process tool calls (resolve names, fill defaults)
                    filled_actual_tool_calls = self._process_tool_calls(
                        predicted_args, registry=use_registry
                    )

                    # Convert to CapturedToolCall objects
                    tool_calls = [
                        CapturedToolCall(name=name, args=args)
                        for name, args in filled_actual_tool_calls
                    ]

                    runs.append(CapturedRun(tool_calls=tool_calls))

                primary_tool_calls = runs[0].tool_calls if runs else []

                return CapturedCase(
                    case_name=case.name,
                    user_message=case.user_message,
                    tool_calls=primary_tool_calls,
                    system_message=case.system_message if include_context else None,
                    additional_messages=case.additional_messages if include_context else None,
                    track_name=track,
                    runs=runs if len(runs) > 1 else [],
                )

        # Capture regular cases (using default registry)
        if self.cases:
            tasks = [capture_case(case) for case in self.cases]
            regular_captured = await asyncio.gather(*tasks)
            all_captured.extend(regular_captured)

        # Capture comparative cases (each track separately)
        if self._comparative_case_builders:
            for builder in self._comparative_case_builders:
                comp_case = builder.build()

                # For each track configured in this comparative case
                for track_name in comp_case.track_configs:
                    if not self._track_manager.has_track(track_name):
                        continue  # Skip missing tracks

                    track_registry = self._track_manager.get_registry(track_name)

                    # Create an EvalCase from the comparative case
                    # Use case-specific rubric if defined, otherwise use suite default
                    case_rubric = comp_case.rubric or self.rubric
                    eval_case = self._create_eval_case(
                        name=comp_case.name,  # Don't embed track in name - use track_name field
                        user_message=comp_case.user_message,
                        system_message=comp_case.system_message,
                        additional_messages=comp_case.additional_messages,
                        expected_tool_calls=[],  # Not needed for capture
                        rubric=case_rubric,
                        critics=[],  # Not needed for capture
                    )

                    captured = await capture_case(
                        eval_case, registry=track_registry, track=track_name
                    )
                    all_captured.append(captured)

        return CaptureResult(
            suite_name=self.name,
            model=model,
            provider=provider,
            captured_cases=list(all_captured),
        )
