"""
Data models for Clew-Bench arena experiments.
"""

from __future__ import annotations

import statistics
from dataclasses import dataclass, field
from enum import Enum


class AgentType(Enum):
    """Identifies which context backend produced a result."""

    CONTROL = "control"  # Standard RAG (vector-only)
    TEST = "test"  # Clew Engine (graph + vector)


@dataclass
class TokenUsage:
    """Token counts from a single LLM call."""

    input_tokens: int = 0
    output_tokens: int = 0

    @property
    def total(self) -> int:
        return self.input_tokens + self.output_tokens


@dataclass
class Judgement:
    """Quality assessment from the Judge LLM."""

    correctness: int  # 1-5
    efficiency: int  # 1-5
    hallucinations: list[str] = field(default_factory=list)
    comments: str = ""

    @property
    def avg_score(self) -> float:
        return (self.correctness + self.efficiency) / 2


@dataclass
class TicketResult:
    """Outcome of running one agent on one ticket."""

    ticket_id: str
    agent_type: AgentType
    code: str
    context: str
    context_tokens: int
    usage: TokenUsage
    judgement: Judgement | None = None
    retrieval_time_s: float = 0.0
    generation_time_s: float = 0.0
    is_trap: bool = False


@dataclass
class AgentStats:
    """Aggregated metrics for one agent across all tickets."""

    agent_type: AgentType
    ticket_results: list[TicketResult] = field(default_factory=list)

    @property
    def total_input_tokens(self) -> int:
        return sum(r.usage.input_tokens for r in self.ticket_results)

    @property
    def total_output_tokens(self) -> int:
        return sum(r.usage.output_tokens for r in self.ticket_results)

    @property
    def quality_scores(self) -> list[float]:
        return [r.judgement.avg_score for r in self.ticket_results if r.judgement]

    @property
    def mean_quality(self) -> float:
        scores = self.quality_scores
        return statistics.mean(scores) if scores else 0.0

    @property
    def hallucination_count(self) -> int:
        return sum(len(r.judgement.hallucinations) for r in self.ticket_results if r.judgement)

    @property
    def hallucination_rate(self) -> float:
        judged = [r for r in self.ticket_results if r.judgement]
        if not judged:
            return 0.0
        return sum(1 for r in judged if r.judgement and r.judgement.hallucinations) / len(judged)

    @property
    def trap_results(self) -> list[TicketResult]:
        return [r for r in self.ticket_results if r.is_trap]

    @property
    def refactor_successes(self) -> int:
        return sum(1 for r in self.trap_results if r.judgement and r.judgement.correctness >= 4)

    @property
    def refactor_total(self) -> int:
        return len(self.trap_results)

    @property
    def context_tokens_per_ticket(self) -> list[int]:
        return [r.context_tokens for r in self.ticket_results]

    def to_dict(self) -> dict:
        """Serialize to dictionary for JSON export."""
        return {
            "agent_type": self.agent_type.value,
            "total_input_tokens": self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
            "mean_quality": round(self.mean_quality, 2),
            "hallucination_count": self.hallucination_count,
            "hallucination_rate": round(self.hallucination_rate, 4),
            "refactor_successes": self.refactor_successes,
            "refactor_total": self.refactor_total,
            "context_tokens_per_ticket": self.context_tokens_per_ticket,
            "tickets": [
                {
                    "ticket_id": r.ticket_id,
                    "context_tokens": r.context_tokens,
                    "input_tokens": r.usage.input_tokens,
                    "output_tokens": r.usage.output_tokens,
                    "retrieval_time_s": round(r.retrieval_time_s, 3),
                    "generation_time_s": round(r.generation_time_s, 3),
                    "is_trap": r.is_trap,
                    "judgement": {
                        "correctness": r.judgement.correctness,
                        "efficiency": r.judgement.efficiency,
                        "avg_score": round(r.judgement.avg_score, 2),
                        "hallucinations": r.judgement.hallucinations,
                        "comments": r.judgement.comments,
                    }
                    if r.judgement
                    else None,
                }
                for r in self.ticket_results
            ],
        }


@dataclass
class Ticket:
    """A single workload ticket parsed from epics.json."""

    id: str
    instruction: str
    success_criteria: str
    trap: bool = False
    notes: str = ""


@dataclass
class Epic:
    """A group of related tickets."""

    epic_id: str
    title: str
    tickets: list[Ticket]


@dataclass
class Workload:
    """Full workload definition loaded from JSON."""

    workload_id: str
    description: str
    track: str
    epics: list[Epic]

    @property
    def all_tickets(self) -> list[Ticket]:
        return [t for e in self.epics for t in e.tickets]

    @property
    def trap_count(self) -> int:
        return sum(1 for t in self.all_tickets if t.trap)

    @classmethod
    def from_dict(cls, data: dict) -> Workload:
        """Parse workload from JSON dictionary."""
        epics = []
        for epic_data in data["epics"]:
            tickets = [
                Ticket(
                    id=t["id"],
                    instruction=t["instruction"],
                    success_criteria=t["success_criteria"],
                    trap=t.get("trap", False),
                    notes=t.get("notes", ""),
                )
                for t in epic_data["tickets"]
            ]
            epics.append(Epic(epic_id=epic_data["epic_id"], title=epic_data["title"], tickets=tickets))
        return cls(
            workload_id=data["workload_id"],
            description=data["description"],
            track=data["track"],
            epics=epics,
        )


@dataclass
class ExperimentResult:
    """Complete result of one arena experiment run."""

    workload_id: str
    control: AgentStats
    test: AgentStats
    trial_number: int = 1

    @property
    def compression_ratio(self) -> float:
        """How many times fewer tokens Clew uses vs RAG."""
        if self.test.total_input_tokens == 0:
            return 0.0
        return self.control.total_input_tokens / self.test.total_input_tokens

    @property
    def compression_pct(self) -> float:
        """Percentage of tokens saved by Clew."""
        if self.control.total_input_tokens == 0:
            return 0.0
        return (1 - self.test.total_input_tokens / self.control.total_input_tokens) * 100

    def to_dict(self) -> dict:
        return {
            "workload_id": self.workload_id,
            "trial_number": self.trial_number,
            "compression_ratio": round(self.compression_ratio, 2),
            "compression_pct": round(self.compression_pct, 1),
            "control": self.control.to_dict(),
            "test": self.test.to_dict(),
        }
