"""Pydantic models for configuration, payloads, and results."""

from __future__ import annotations

from enum import Enum
from typing import Any

from pydantic import BaseModel, Field


# --- Target Configuration ---


class AuthType(str, Enum):
    none = "none"
    bearer = "bearer"
    header = "header"
    basic = "basic"


class AuthConfig(BaseModel):
    type: AuthType = AuthType.none
    token: str | None = None
    header_name: str | None = None
    header_value: str | None = None
    username: str | None = None
    password: str | None = None


class RequestTemplate(BaseModel):
    method: str = "POST"
    headers: dict[str, str] = Field(default_factory=dict)
    body_template: dict[str, Any] = Field(default_factory=dict)


class TargetConfig(BaseModel):
    name: str = "target"
    endpoint: str
    auth: AuthConfig = Field(default_factory=AuthConfig)
    request: RequestTemplate = Field(default_factory=RequestTemplate)
    response_path: str = "choices[0].message.content"
    rate_limit: float = 10.0
    timeout: float = 30.0


# --- Attack Payloads ---


class AttackCategory(str, Enum):
    prompt_injection = "prompt_injection"
    data_exfiltration = "data_exfiltration"
    tool_abuse = "tool_abuse"
    jailbreak = "jailbreak"


class Severity(str, Enum):
    info = "info"
    low = "low"
    medium = "medium"
    high = "high"
    critical = "critical"


SEVERITY_ORDER = {"info": 0, "low": 1, "medium": 2, "high": 3, "critical": 4}


class DetectionIndicator(BaseModel):
    pattern: str
    confidence: float = 0.5


class TurnPrompt(BaseModel):
    """A single turn in a multi-turn attack."""
    role: str = "user"
    content: str
    detect: bool = False


class Payload(BaseModel):
    id: str
    name: str
    category: AttackCategory
    subcategory: str = ""
    severity: Severity = Severity.medium
    description: str = ""
    prompts: list[str] = Field(default_factory=list)
    turns: list[list[TurnPrompt]] = Field(default_factory=list)
    detection: list[DetectionIndicator] = Field(default_factory=list)
    tags: list[str] = Field(default_factory=list)
    owasp: list[str] = Field(default_factory=list)  # e.g. ["LLM01", "LLM07"]


# --- Chains ---


class ChainStep(BaseModel):
    payload_id: str
    prompt_index: int = 0  # Which prompt from the payload to use


class Chain(BaseModel):
    id: str
    name: str
    description: str = ""
    steps: list[ChainStep]


# --- Results ---


class AttackResult(BaseModel):
    payload_id: str
    prompt: str
    response: str
    success: bool = False
    confidence: float = 0.0
    matched_indicators: list[str] = Field(default_factory=list)
    duration_ms: float = 0.0
    error: str | None = None
    severity: Severity = Severity.medium
    owasp: list[str] = Field(default_factory=list)


class ScanResult(BaseModel):
    target: str
    timestamp: str = ""
    total_attacks: int = 0
    successful_attacks: int = 0
    results: list[AttackResult] = Field(default_factory=list)
