from typing import Any, Literal

from pydantic import BaseModel, Field


class ChatRequest(BaseModel):
    message: str
    scenario_id: str = "prompt_injection"
    sample_id: str | None = None
    mode: Literal["vulnerable", "secure"] = "secure"
    task: str | None = None


class SecurityPipelineStep(BaseModel):
    id: str
    name: str
    status: Literal["passed", "blocked", "skipped", "warning"]
    enabled: bool
    summary: str
    detail: str | None = None


class ChatResult(BaseModel):
    output: str
    attack_success: bool
    consequences: list[str] = Field(default_factory=list)
    blocked_by: str | None = None
    defense_explanation: str | None = None
    risk_level: Literal["low", "medium", "high", "critical"] = "low"
    audit_entries: list[dict[str, Any]] = Field(default_factory=list)
    security_pipeline: list[SecurityPipelineStep] = Field(default_factory=list)


class AttackMeta(BaseModel):
    scenario: str
    scenario_name: str
    sample_id: str | None = None
    sample_name: str | None = None
    intent: str | None = None
    expected_consequences: list[str] = Field(default_factory=list)


class CompareResponse(BaseModel):
    attack_meta: AttackMeta
    vulnerable: ChatResult
    secure: ChatResult


class SecurityConfigUpdate(BaseModel):
    enable_input_isolation: bool | None = None
    enable_injection_detection: bool | None = None
    enable_jailbreak_detection: bool | None = None
    enable_output_filter: bool | None = None
    task_whitelist: list[str] | None = None
    refusal_message: str | None = None
    audit_enabled: bool | None = None
    llm_provider: str | None = None


class AuditLogEntry(BaseModel):
    timestamp: str
    scenario: str
    mode: str
    message_preview: str
    attack_success: bool
    blocked_by: str | None = None
    consequences: list[str] = Field(default_factory=list)