"""Core data models for rewardprobe."""

from __future__ import annotations

import random
from enum import Enum
from typing import Any, Callable, Optional

from pydantic import BaseModel, Field, model_validator


class Severity(str, Enum):
    CRITICAL = "critical"
    WARNING = "warning"
    INFO = "info"


class Finding(BaseModel):
    """A single issue found by an attack probe."""

    attack: str
    family: str
    severity: Severity
    input: str = ""
    expected_reward: Optional[float] = None
    actual_reward: Optional[float] = None
    affected_tasks: list[int] = Field(default_factory=list)
    error: Optional[str] = None
    recommendation: str = ""


class NormalizedRewardFn(BaseModel):
    """A reward function normalized to a simple sync interface."""

    model_config = {"arbitrary_types_allowed": True}

    name: str
    fn: Callable[[str, dict], float]
    weight: float = 1.0


class NormalizedEnvironment(BaseModel):
    """What every adapter produces — the universal interface for attacks."""

    model_config = {"arbitrary_types_allowed": True}

    name: str
    reward_fns: list[NormalizedRewardFn]
    tasks: list[dict[str, Any]]
    parser: Optional[Any] = None
    metadata: dict[str, Any] = Field(default_factory=dict)

    @model_validator(mode="after")
    def _validate(self):
        if not self.reward_fns:
            raise ValueError("NormalizedEnvironment requires at least one reward function")
        if not self.tasks:
            raise ValueError("NormalizedEnvironment requires at least one task")
        return self

    def sample_tasks(self, n: int = 10, seed: int = 42) -> list[dict[str, Any]]:
        """Deterministically sample n tasks for reproducibility."""
        if len(self.tasks) <= n:
            return list(self.tasks)
        rng = random.Random(seed)
        return rng.sample(self.tasks, n)

    def compute_rubric_score(self, completion: str, task: dict) -> float:
        """Compute weighted sum of all reward functions (the rubric score)."""
        total_weight = sum(fn.weight for fn in self.reward_fns)
        if total_weight == 0:
            return 0.0
        weighted_sum = sum(fn.fn(completion, task) * fn.weight for fn in self.reward_fns)
        return weighted_sum / total_weight


class FamilyResult(BaseModel):
    """Results from running one attack family."""

    name: str
    display_name: str
    total: int
    passed: int
    findings: list[Finding] = Field(default_factory=list)

    @property
    def all_passed(self) -> bool:
        return self.passed == self.total


class Report(BaseModel):
    """Full probe report."""

    environment: str
    rewardprobe_version: str
    timestamp: str
    passed: bool
    families: list[FamilyResult] = Field(default_factory=list)

    @property
    def findings(self) -> list[Finding]:
        return [f for family in self.families for f in family.findings]

    @property
    def critical_count(self) -> int:
        return sum(1 for f in self.findings if f.severity == Severity.CRITICAL)

    @property
    def warning_count(self) -> int:
        return sum(1 for f in self.findings if f.severity == Severity.WARNING)

    @property
    def total_attacks(self) -> int:
        return sum(f.total for f in self.families)

    @property
    def total_passed(self) -> int:
        return sum(f.passed for f in self.families)

    def to_dict(self) -> dict:
        from rewardprobe.report import report_to_dict

        return report_to_dict(self)

    def to_json(self, path: str) -> None:
        from rewardprobe.report import report_to_json

        report_to_json(self, path)
