"""
Arena Harness: orchestrates side-by-side experiments.

Runs the same workload through both context engines (Clew vs Standard RAG),
sends retrieved context to the coding LLM, then has the judge LLM score
each submission. Collects all metrics for the final report.
"""

from __future__ import annotations

import json
import logging
import time
from pathlib import Path

from .coding_llm import CodingLLM
from .engines.base import ContextEngine
from .judge import JudgeLLM
from .models import AgentStats, AgentType, ExperimentResult, TicketResult, Workload

logger = logging.getLogger("clew.bench.arena")


class ArenaHarness:
    """
    Main experiment orchestrator.

    For each ticket in the workload:
    1. Retrieve context via the control engine (Standard RAG)
    2. Retrieve context via the test engine (Clew)
    3. Generate code with the coding LLM for both
    4. Judge both submissions
    5. Collect metrics
    """

    def __init__(
        self,
        workload: Workload,
        clew_engine: ContextEngine,
        standard_engine: ContextEngine,
        coding_llm: CodingLLM,
        judge: JudgeLLM,
    ):
        self.workload = workload
        self.clew_engine = clew_engine
        self.standard_engine = standard_engine
        self.coding_llm = coding_llm
        self.judge = judge

    async def _run_agent(
        self,
        agent_type: AgentType,
        engine: ContextEngine,
        ticket_id: str,
        instruction: str,
        success_criteria: str,
        is_trap: bool,
    ) -> TicketResult:
        """Run one agent (context retrieval + code generation) on one ticket."""
        # 1. Retrieve context
        retrieval_start = time.monotonic()
        ctx = await engine.query(instruction)
        retrieval_time = time.monotonic() - retrieval_start

        # 2. Generate code
        gen_start = time.monotonic()
        code, usage = await self.coding_llm.generate(instruction, success_criteria, ctx.context)
        gen_time = time.monotonic() - gen_start

        logger.info(
            "[%s] %s | context=%d tokens | gen=%d/%d tokens | retrieval=%.2fs | gen=%.2fs",
            agent_type.value,
            ticket_id,
            ctx.token_count,
            usage.input_tokens,
            usage.output_tokens,
            retrieval_time,
            gen_time,
        )

        return TicketResult(
            ticket_id=ticket_id,
            agent_type=agent_type,
            code=code,
            context=ctx.context,
            context_tokens=ctx.token_count,
            usage=usage,
            retrieval_time_s=retrieval_time,
            generation_time_s=gen_time,
            is_trap=is_trap,
        )

    async def _judge_result(self, result: TicketResult, instruction: str, success_criteria: str) -> None:
        """Have the judge LLM evaluate a ticket result."""
        judgement = await self.judge.evaluate(instruction, success_criteria, result.code)
        result.judgement = judgement

        level = logging.WARNING if judgement.hallucinations else logging.INFO
        logger.log(
            level,
            "[%s] %s | judge: correctness=%d efficiency=%d hallucinations=%d",
            result.agent_type.value,
            result.ticket_id,
            judgement.correctness,
            judgement.efficiency,
            len(judgement.hallucinations),
        )

    async def run_experiment(self, trial_number: int = 1) -> ExperimentResult:
        """
        Run a complete experiment across all tickets in the workload.

        Args:
            trial_number: Trial identifier for multi-run experiments.

        Returns:
            ExperimentResult with aggregated metrics for both agents.
        """
        control_stats = AgentStats(agent_type=AgentType.CONTROL)
        test_stats = AgentStats(agent_type=AgentType.TEST)

        for epic in self.workload.epics:
            logger.info("--- Epic: %s (%s) ---", epic.title, epic.epic_id)

            for ticket in epic.tickets:
                logger.info("Ticket %s: %s", ticket.id, ticket.instruction[:80])

                # Run control (Standard RAG)
                result_c = await self._run_agent(
                    AgentType.CONTROL,
                    self.standard_engine,
                    ticket.id,
                    ticket.instruction,
                    ticket.success_criteria,
                    ticket.trap,
                )
                await self._judge_result(result_c, ticket.instruction, ticket.success_criteria)
                control_stats.ticket_results.append(result_c)

                # Run test (Clew)
                result_t = await self._run_agent(
                    AgentType.TEST,
                    self.clew_engine,
                    ticket.id,
                    ticket.instruction,
                    ticket.success_criteria,
                    ticket.trap,
                )
                await self._judge_result(result_t, ticket.instruction, ticket.success_criteria)
                test_stats.ticket_results.append(result_t)

        return ExperimentResult(
            workload_id=self.workload.workload_id,
            control=control_stats,
            test=test_stats,
            trial_number=trial_number,
        )


def load_workload(path: str | Path) -> Workload:
    """Load a workload definition from a JSON file."""
    with open(path) as f:
        data = json.load(f)
    return Workload.from_dict(data)
