from __future__ import annotations

from typing import Literal

from pydantic import BaseModel, Field

ClauseType = Literal[
    "non_compete",
    "arbitration",
    "ip_assignment",
    "auto_renewal",
    "termination",
    "liability",
    "data_collection",
    "payment",
    "confidentiality",
    "indemnification",
    "warranty",
    "governing_law",
    "other",
]

Severity = Literal["low", "medium", "high", "critical"]

DocType = Literal[
    "employment",
    "freelance",
    "saas_tos",
    "privacy_policy",
    "rental",
    "vendor",
    "nda",
    "other",
]

AgentName = Literal["extractor", "risk_analyst", "devil_advocate", "legal_context", "negotiator"]

SEVERITY_RANK: dict[Severity, int] = {"low": 0, "medium": 1, "high": 2, "critical": 3}


class NegotiationSuggestion(BaseModel):
    rewrite: str
    talking_points: list[str]
    walk_away_threshold: str | None = None


class DocumentEntities(BaseModel):
    parties: list[str] = Field(default_factory=list)
    dates: list[str] = Field(default_factory=list)
    monetary_amounts: list[str] = Field(default_factory=list)
    durations: list[str] = Field(default_factory=list)
    jurisdictions: list[str] = Field(default_factory=list)
    key_obligations: list[str] = Field(default_factory=list)


class Clause(BaseModel):
    id: str
    type: ClauseType
    text: str
    span: tuple[int, int]
    heading: str | None = None
    severity: Severity = "low"
    risk_score: float = Field(default=0.0, ge=0.0, le=100.0)
    rationale: str = ""
    worst_case: str | None = None
    benchmark_delta: str | None = None
    negotiation: NegotiationSuggestion | None = None
    # Explainable AI fields
    confidence_score: float = Field(default=50.0, ge=0.0, le=100.0)
    legal_principles: list[str] = Field(default_factory=list)
    statutory_refs: list[str] = Field(default_factory=list)
    # RAG / semantic similarity
    similarity_score: float | None = None
    # NLP entity extraction (clause-level)
    entities: DocumentEntities | None = None


class AgentMessage(BaseModel):
    agent: AgentName
    clause_id: str | None = None
    content: str
    timestamp: float


class BenchmarkRef(BaseModel):
    id: str
    doc_type: DocType
    clause_type: ClauseType
    severity_baseline: Severity
    text: str
    notes: str = ""


class BenchmarkMatch(BaseModel):
    ref: BenchmarkRef
    similarity: float = 0.0


class RiskReport(BaseModel):
    job_id: str
    doc_type: DocType
    overall_risk: float = Field(default=0.0, ge=0.0, le=100.0)
    clauses: list[Clause] = Field(default_factory=list)
    summary: str = ""
    agent_trace: list[AgentMessage] = Field(default_factory=list)
    document_entities: DocumentEntities | None = None
    tech_stack: list[str] = Field(default_factory=list)


class AnalyzeResponse(BaseModel):
    job_id: str
    stream_url: str


class HealthResponse(BaseModel):
    status: Literal["ok"]
    version: str
    model: str


Clause.model_rebuild()