"""BenchmarkRunner — orchestrates comparative control/treatment runs.

Control = raw LLM generation (single prompt, no pipeline).
Treatment = full AgentGuard pipeline (skeleton → contracts → wiring → logic +
            validation + self-challenge).

For each benchmark spec the runner:
  1. Generates code via the control path (raw LLM).
  2. Generates code via the treatment path (Pipeline).
  3. Evaluates both outputs with the enterprise + operational evaluator.
  4. Records per-complexity results and computes deltas.
  5. Produces a signed BenchmarkReport.

The runner accepts an ``LLMProvider`` (or model string) directly, so callers
who already have an LLM — MCP agents, CrewAI/LangGraph/OpenHands — never
need to supply a separate API key.
"""

from __future__ import annotations

import logging
import time as _time
from collections.abc import Awaitable, Callable
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, cast

from agentguard.benchmark.types import (
    BenchmarkConfig,
    BenchmarkCriterion,
    BenchmarkReport,
    BenchmarkSpec,
    Complexity,
    ComplexityRun,
    ReadinessScore,
    RunResult,
)

if TYPE_CHECKING:
    from agentguard.archetypes.base import Archetype

_log = logging.getLogger(__name__)

# Type alias for the progress callback.
_ProgressCB = Callable[..., Awaitable[None]]

# Type alias for the resolved evaluator:
#   (spec, files, enterprise_threshold, operational_threshold, llm)
#   -> (enterprise_score, operational_score)
_EvaluatorFn = Callable[
    ...,
    Awaitable[tuple[ReadinessScore, ReadinessScore]],
]

# Minimal system prompt for the control (raw LLM) run.
_CONTROL_SYSTEM = (
    "You are an expert software developer. Generate production-quality code "
    "for the following specification. Return the code as a set of files. "
    "Use the format:\n\n"
    "```<filepath>\n<content>\n```\n\n"
    "for each file. Include all necessary imports and boilerplate."
)


# ── Evaluator wrappers ────────────────────────────────────────────────────────
# Both adapters share the same async signature used by _run_control/_run_treatment:
#   async (spec, files, ent_thresh, ops_thresh, llm) -> (enterprise, operational)


def _wrap_profile(profile: Any) -> _EvaluatorFn:
    """Wrap a synchronous profile evaluator as an async callable."""
    async def _eval(
        spec: str,
        files: dict[str, str],
        enterprise_threshold: float,
        operational_threshold: float,
        llm: object,
    ) -> tuple[ReadinessScore, ReadinessScore]:
        return cast(
            "tuple[ReadinessScore, ReadinessScore]",
            profile.evaluate(spec, files, enterprise_threshold, operational_threshold),
        )
    return _eval


def _wrap_criteria(evaluator: Any) -> _EvaluatorFn:
    """Wrap a CriteriaBasedEvaluator as the standard evaluator signature."""
    async def _eval(
        spec: str,
        files: dict[str, str],
        enterprise_threshold: float,
        operational_threshold: float,
        llm: object,
    ) -> tuple[ReadinessScore, ReadinessScore]:
        return cast(
            "tuple[ReadinessScore, ReadinessScore]",
            await evaluator.evaluate(spec, files, llm),
        )
    return _eval


def _extract_inline_specs(archetype: Archetype) -> list[BenchmarkSpec]:
    """Extract inline benchmark specs from the archetype's benchmark config."""
    bench = getattr(archetype, "benchmark_config", None)
    if bench is None or not bench.specs:
        return []

    category = getattr(archetype, "id", "general")
    result: list[BenchmarkSpec] = []
    for complexity_str, spec_text in bench.specs.items():
        try:
            complexity = Complexity(complexity_str)
            result.append(BenchmarkSpec(complexity=complexity, spec=spec_text, category=category))
        except ValueError:
            _log.warning("Unknown complexity '%s' in inline specs — skipped", complexity_str)
    return result


class BenchmarkRunner:
    """Run a comparative benchmark (control vs treatment) for an archetype.

    Usage::

        from agentguard.benchmark import BenchmarkRunner
        from agentguard.benchmark.catalog import get_default_specs

        config = BenchmarkConfig(
            model="anthropic/claude-sonnet-4-20250514",
            specs=get_default_specs("backend"),
        )
        runner = BenchmarkRunner(archetype="api_backend", config=config)
        report = await runner.run()
        report.sign(secret="my-signing-key")
        print(report.to_json())
    """

    def __init__(
        self,
        archetype: Archetype | str,
        config: BenchmarkConfig,
        *,
        llm: Any | str | None = None,
        signing_secret: str = "",
    ) -> None:
        """Initialise a benchmark runner.

        Args:
            archetype: Archetype instance or built-in name.
            config: Benchmark configuration (specs, thresholds, budget).
            llm: LLM to use for control + treatment runs.  Accepts an
                already-instantiated ``LLMProvider`` (no extra API key
                needed) **or** a model string like
                ``"anthropic/claude-sonnet-4-20250514"`` which will be
                resolved via ``create_llm_provider``.  If *None*, falls
                back to ``config.model`` (CLI convenience).
            signing_secret: HMAC secret for report signing.
        """
        from agentguard.archetypes.registry import get_archetype_registry
        from agentguard.benchmark._compat import create_llm_provider

        # Resolve archetype
        if isinstance(archetype, str):
            registry = get_archetype_registry()
            self._archetype: Archetype = registry.get(archetype)
        else:
            self._archetype = archetype

        # Populate specs from archetype inline benchmark config when caller
        # hasn't provided any.  This lets archetype authors ship their own
        # benchmark specs inside the YAML without requiring 5 complexity levels.
        if not config.specs:
            inline_specs = _extract_inline_specs(self._archetype)
            if inline_specs:
                config.specs = inline_specs
                config.require_all_complexities = False
                _log.info(
                    "Using %d inline benchmark spec(s) from archetype '%s'",
                    len(inline_specs),
                    self._archetype.id,
                )

        # Validate config
        errors = config.validate()
        if errors:
            raise ValueError(f"Invalid benchmark config: {'; '.join(errors)}")

        self._config = config

        # Resolve LLM: direct provider > model string > config.model
        if isinstance(llm, str):
            self._llm: Any = create_llm_provider(llm)
            # Store model label for the report
            if not config.model:
                config.model = llm
        elif llm is not None:
            self._llm = llm
            if not config.model:
                config.model = getattr(llm, "model", "unknown")
        elif config.model:
            self._llm = create_llm_provider(config.model)
        else:
            raise ValueError(
                "No LLM provided.  Pass an LLMProvider instance, a model "
                "string, or set config.model."
            )

        self._secret = signing_secret
        self._total_cost: float = 0.0

        # Resolve which evaluator to use for this archetype's output type
        self._evaluator = self._resolve_evaluator()

    async def run(self, *, progress_callback: _ProgressCB | None = None) -> BenchmarkReport:
        """Execute the full benchmark across all configured specs.

        Args:
            progress_callback: Optional async callable(complexity, step, detail)
                for reporting progress to callers (CLI, MCP, etc.).

        Returns:
            A BenchmarkReport with all run results and aggregates.
        """
        runs: list[ComplexityRun] = []
        self._total_cost = 0.0

        for i, spec in enumerate(self._config.specs):
            _log.info(
                "Benchmark [%d/%d] complexity=%s spec=%s",
                i + 1, len(self._config.specs),
                spec.complexity.value,
                spec.spec[:60],
            )

            if progress_callback:
                await progress_callback(
                    spec.complexity.value,
                    "starting",
                    f"Running {spec.complexity.value} benchmark ({i + 1}/{len(self._config.specs)})",
                )

            # Budget guard
            if self._total_cost >= self._config.budget_ceiling_usd:
                _log.warning(
                    "Budget ceiling (%.2f USD) reached after %d runs — stopping",
                    self._config.budget_ceiling_usd,
                    len(runs),
                )
                break

            # ── Control run ────────────────────────────────────
            if progress_callback:
                await progress_callback(spec.complexity.value, "control", "Raw LLM generation (no AgentGuard)")
            control = await self._run_control(spec.spec)

            # ── Treatment run ──────────────────────────────────
            if progress_callback:
                await progress_callback(spec.complexity.value, "treatment", "AgentGuard pipeline generation")
            treatment = await self._run_treatment(spec.spec)

            complexity_run = ComplexityRun(
                complexity=spec.complexity,
                spec=spec.spec,
                control=control,
                treatment=treatment,
            )
            runs.append(complexity_run)

            self._total_cost += control.cost_usd + treatment.cost_usd

            _log.info(
                "  → control=%.3f  treatment=%.3f  Δ=%.3f",
                control.combined_score,
                treatment.combined_score,
                complexity_run.improvement,
            )

        # Build report
        report = self._build_report(runs)

        if progress_callback:
            await progress_callback(
                "all", "complete",
                f"Benchmark complete: {'PASSED' if report.overall_passed else 'FAILED'} "
                f"(Δ={report.improvement_avg:.3f}, cost=${report.total_cost_usd:.4f})",
            )

        return report

    # ══════════════════════════════════════════════════════════
    #  Control run — raw LLM, no pipeline
    # ══════════════════════════════════════════════════════════

    async def _run_control(self, spec: str) -> RunResult:
        """Generate code with a raw LLM call (no AgentGuard)."""
        from agentguard.types import Message

        t0 = _time.perf_counter()
        try:
            response = await self._llm.generate(
                messages=[
                    Message(role="system", content=_CONTROL_SYSTEM),
                    Message(role="user", content=spec),
                ],
            )
            duration_ms = int((_time.perf_counter() - t0) * 1000)

            # Parse file blocks from the LLM response
            files = _parse_file_blocks(response.content)
            if not files:
                # Treat the whole response as a single file
                files = {"main.py": response.content}

            # Evaluate using the resolved evaluator for this archetype type
            enterprise, operational = await self._evaluator(
                spec, files,
                self._config.enterprise_threshold,
                self._config.operational_threshold,
                self._llm,
            )

            total_lines = sum(c.count("\n") + 1 for c in files.values())
            cost = float(response.cost.total_cost)

            return RunResult(
                enterprise=enterprise,
                operational=operational,
                files_generated=len(files),
                total_lines=total_lines,
                total_tokens=response.tokens.total_tokens,
                cost_usd=cost,
                duration_ms=duration_ms,
            )
        except Exception as exc:
            duration_ms = int((_time.perf_counter() - t0) * 1000)
            _log.error("Control run failed: %s", exc)
            return _error_result(str(exc), duration_ms)

    # ══════════════════════════════════════════════════════════
    #  Treatment run — full AgentGuard pipeline
    # ══════════════════════════════════════════════════════════

    async def _run_treatment(self, spec: str) -> RunResult:
        """Generate code through the AgentGuard pipeline.

        Note: In v0.6.0 the internal Pipeline was removed. Treatment runs
        now return an error result. Use the agent-native benchmark tools
        (benchmark + benchmark_evaluate) instead.
        """
        t0 = _time.perf_counter()
        duration_ms = int((_time.perf_counter() - t0) * 1000)
        _log.warning(
            "Treatment run skipped: Pipeline class was removed in v0.6.0. "
            "Use agent-native benchmark tools (benchmark + benchmark_evaluate) instead."
        )
        return _error_result(
            "Pipeline removed in v0.6.0 — use agent-native benchmark tools",
            duration_ms,
        )

    # ══════════════════════════════════════════════════════════
    #  Evaluator resolution
    # ══════════════════════════════════════════════════════════

    def _resolve_evaluator(self) -> _EvaluatorFn:
        """Choose the evaluation function for this archetype's output type.

        Resolution chain (first match wins):

        1. ``archetype.benchmark_config.profile`` → named profile in registry.
        2. ``archetype.benchmark_config.criteria`` → ``CriteriaBasedEvaluator``
           (LLM-judge, one pass per criterion).
        3. Fallback → ``"generic"`` profile (content, structure, spec coverage).
        """
        from agentguard.benchmark.profiles import get_profile

        bench = getattr(self._archetype, "benchmark_config", None)

        # ── 1. Named profile ──────────────────────────────────
        if bench and bench.profile:
            profile = get_profile(bench.profile)
            if profile is not None:
                _log.info("Benchmark evaluator: named profile '%s'", bench.profile)
                return _wrap_profile(profile)
            _log.warning(
                "Profile '%s' not found in registry — falling back to criteria/generic",
                bench.profile,
            )

        # ── 2. Inline criteria → LLM judge ───────────────────
        if bench and bench.criteria:
            from agentguard.benchmark.criteria_evaluator import CriteriaBasedEvaluator

            criteria = [
                BenchmarkCriterion(
                    name=c.name,
                    description=c.description,
                    rubric=c.rubric,
                    weight=c.weight,
                )
                for c in bench.criteria
            ]
            evaluator = CriteriaBasedEvaluator(
                criteria,
                threshold=self._config.enterprise_threshold,
            )
            _log.info(
                "Benchmark evaluator: criteria-based LLM judge (%d criteria)",
                len(criteria),
            )
            return _wrap_criteria(evaluator)

        # ── 3. Generic fallback ───────────────────────────────
        _log.info("Benchmark evaluator: generic profile (fallback)")
        generic = get_profile("generic")
        if generic is None:
            # Profiles module not imported yet — trigger registration
            from agentguard.benchmark.profiles import builtin as _b  # noqa: F401
            generic = get_profile("generic")
        return _wrap_profile(generic)

    # ══════════════════════════════════════════════════════════
    #  Report building
    # ══════════════════════════════════════════════════════════

    def _build_report(self, runs: list[ComplexityRun]) -> BenchmarkReport:
        """Build and optionally sign the benchmark report."""
        # Compute archetype hash
        try:
            import yaml  # noqa: F401

            from agentguard.archetypes.schema import compute_content_hash  # noqa: F401

            # Hash the archetype YAML if available
            arch_hash = getattr(self._archetype, "_yaml_content_hash", "")
            if not arch_hash:
                arch_hash = ""
        except Exception:
            arch_hash = ""

        report = BenchmarkReport(
            archetype_id=self._archetype.id,
            archetype_hash=arch_hash,
            model=self._config.model,
            runs=runs,
            created_at=datetime.now(UTC).isoformat(),
        )
        report.compute_aggregates(
            enterprise_threshold=self._config.enterprise_threshold,
            operational_threshold=self._config.operational_threshold,
            improvement_threshold=self._config.improvement_threshold,
        )

        # Additional pass condition: all levels must complete
        report.overall_passed = (
            report.overall_passed
            and len(runs) == len(self._config.specs)
        )

        # Sign if secret provided
        if self._secret:
            report.sign(self._secret)

        return report


# ══════════════════════════════════════════════════════════════════
#  Helpers
# ══════════════════════════════════════════════════════════════════


def _error_result(error: str, duration_ms: int) -> RunResult:
    """Create a zero-score RunResult for an errored run."""
    empty_enterprise = ReadinessScore(
        category="enterprise", overall_score=0.0, passed=False, dimensions=[],
    )
    empty_operational = ReadinessScore(
        category="operational", overall_score=0.0, passed=False, dimensions=[],
    )
    return RunResult(
        enterprise=empty_enterprise,
        operational=empty_operational,
        duration_ms=duration_ms,
        error=error,
    )


def _parse_file_blocks(content: str) -> dict[str, str]:
    """Parse ```filepath ... ``` code blocks from LLM output.

    Supports formats:
      ```python filepath
      ```filepath
      ```<filepath>
    """
    import re

    pattern = re.compile(
        r"```(?:\w+\s+)?<?([^\s>`]+)>?\s*\n(.*?)```",
        re.DOTALL,
    )
    files: dict[str, str] = {}
    for match in pattern.finditer(content):
        fpath = match.group(1).strip()
        fcontent = match.group(2)
        if fpath and not fpath.startswith("#"):
            files[fpath] = fcontent
    return files
