"""Shared data models for SecAgent pipeline."""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Any


class Severity(str, Enum):
    """Vulnerability severity levels."""

    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    INFO = "info"
    UNKNOWN = "unknown"


class VerificationStatus(str, Enum):
    """Whether a finding is a true or false positive."""

    TP = "TP"
    FP = "FP"


# ── Token pricing ─────────────────────────────────────────────────────────────
# (input_per_million_usd, output_per_million_usd)
_MODEL_PRICING: dict[str, tuple[float, float]] = {
    "grok-4-1-fast-reasoning":     (0.20, 0.50),
    "grok-4-1-fast-non-reasoning": (0.20, 0.50),  # update if xAI changes pricing
    "grok-3":                      (3.00, 15.00),
    "grok-2":                      (2.00, 10.00),
    "grok-beta":                   (5.00, 15.00),
}
_DEFAULT_PRICING: tuple[float, float] = (5.00, 15.00)  # conservative fallback


@dataclass
class TokenUsage:
    """Token counts for one or more LLM API calls.

    ``reasoning_tokens`` is a subset of ``completion_tokens`` used by
    reasoning models (e.g. grok-4-1-fast-reasoning) for internal chain-of-
    thought.  It is already included in ``completion_tokens`` for billing.
    """

    prompt_tokens: int = 0
    completion_tokens: int = 0
    reasoning_tokens: int = 0  # subset of completion_tokens
    calls: int = 0

    def add(self, prompt: int, completion: int, reasoning: int = 0) -> None:
        """Accumulate tokens from a single API call."""
        self.prompt_tokens += prompt
        self.completion_tokens += completion
        self.reasoning_tokens += reasoning
        self.calls += 1

    def merge(self, other: TokenUsage) -> None:
        """Merge another TokenUsage into this one in-place."""
        self.prompt_tokens += other.prompt_tokens
        self.completion_tokens += other.completion_tokens
        self.reasoning_tokens += other.reasoning_tokens
        self.calls += other.calls

    @property
    def total_tokens(self) -> int:
        return self.prompt_tokens + self.completion_tokens

    def cost_usd(self, model: str) -> float:
        """Return estimated cost in USD for the accumulated tokens."""
        in_price, out_price = _MODEL_PRICING.get(model, _DEFAULT_PRICING)
        return (
            self.prompt_tokens * in_price + self.completion_tokens * out_price
        ) / 1_000_000


@dataclass
class Vulnerability:
    """Unified vulnerability representation across all scanners."""

    scanner: str
    file: str
    line: int
    severity: Severity
    vuln_type: str
    description: str
    raw: dict[str, Any] = field(default_factory=dict)
    # Populated after verification
    verification: VerificationStatus | None = None

    @property
    def id(self) -> str:
        """Short human-readable identifier."""
        return f"{self.scanner}:{self.file}:{self.line}"


@dataclass
class FixResult:
    """Output produced by the Fixer Agent."""

    file: str
    original_snippet: str
    patched_snippet: str
    reason: str
    vulnerability: Vulnerability


@dataclass
class ValidationResult:
    """Output produced by the Validator Agent."""

    valid: bool
    reason: str
    fix: FixResult


@dataclass
class ScanResult:
    """Aggregated result of a scan-only run (stages 1-4)."""

    target: str
    total_findings: int = 0
    vulnerabilities: list[Vulnerability] = field(default_factory=list)
    scan_file: str | None = None  # path to saved .vulns.json
    stage_usage: dict[str, TokenUsage] = field(default_factory=dict)
    stage_models: dict[str, str] = field(default_factory=dict)  # stage → model name


@dataclass
class SkippedFix:
    """A TP vulnerability whose fix could not be applied."""

    vulnerability: Vulnerability
    reason: str  # e.g. "no fix generated", "validation failed", "patch failed"


@dataclass
class PipelineResult:
    """Aggregated result of a full pipeline run."""

    target: str
    total_findings: int = 0
    true_positives: int = 0
    false_positives: int = 0
    fixes_applied: int = 0
    fixes_validated: int = 0
    vulnerabilities: list[Vulnerability] = field(default_factory=list)
    fixes: list[FixResult] = field(default_factory=list)
    validations: list[ValidationResult] = field(default_factory=list)
    skipped_fixes: list[SkippedFix] = field(default_factory=list)
    report_path: str | None = None
    stage_usage: dict[str, TokenUsage] = field(default_factory=dict)
    stage_models: dict[str, str] = field(default_factory=dict)  # stage → model name
