"""Typed capability requests / responses.

Business services construct these Pydantic models and call methods on the
AIGateway. They never know which concrete model or provider was chosen.
"""

from __future__ import annotations

from typing import Any, Optional, Sequence
from uuid import UUID

from pydantic import BaseModel, Field


class TraceContext(BaseModel):
    """Context propagated through every call for audit / logging."""

    tenant_id: UUID
    project_id: Optional[UUID] = None
    user_id: Optional[UUID] = None
    job_id: Optional[UUID] = None
    correlation_id: Optional[str] = None


class CapabilityRequest(BaseModel):
    """Base marker for every capability request supported by the gateway."""

    trace: TraceContext
    preferred_model: Optional[str] = Field(
        None, max_length=128, description="Hint only. Gateway's routing decision is authoritative."
    )
    priority: str = "normal"


class CapabilityResponse(BaseModel):
    """Every response carries at least trace + usage."""

    provider: str
    model_used: str
    input_tokens: Optional[int] = None
    output_tokens: Optional[int] = None
    duration_ms: Optional[int] = None
    estimated_cost: Optional[float] = None
    execution_id: UUID
    raw: Optional[dict[str, Any]] = Field(
        default=None, description="Raw provider output, preserved for debugging."
    )


# ============================================================================
# Individual capabilities
# ============================================================================


class ChatMessage(BaseModel):
    role: str = Field(..., pattern="^(system|user|assistant|tool)$")
    content: str


class ChatRequest(CapabilityRequest):
    messages: Sequence[ChatMessage]
    temperature: float = Field(0.2, ge=0.0, le=2.0)
    max_tokens: Optional[int] = Field(None, ge=1)
    context_summary: Optional[str] = Field(
        None, description="Short plaintext summary of the Context Engine assembly."
    )


class ChatResponse(CapabilityResponse):
    reply: str
    references: list[str] = Field(default_factory=list)


class CodeAnalysisRequest(CapabilityRequest):
    file_path: Optional[str] = None
    language: Optional[str] = None
    code_snippet: str = Field(..., min_length=1)
    analysis_types: Sequence[str] = Field(
        default_factory=lambda: ["bugs", "complexity", "security", "style"],
        description="Which kinds of analyses to run.",
    )


class Finding(BaseModel):
    severity: str = Field(..., pattern="^(info|low|medium|high|critical)$")
    kind: str
    message: str
    line_start: Optional[int] = None
    line_end: Optional[int] = None
    suggestion: Optional[str] = None


class CodeAnalysisResult(CapabilityResponse):
    summary: str
    findings: list[Finding] = Field(default_factory=list)


class TestGenerationRequest(CapabilityRequest):
    file_path: Optional[str] = None
    language: str = "python"
    test_framework: str = "pytest"
    code_snippet: str = Field(..., min_length=1)
    goal: str = "Generate unit tests covering happy path + main failure modes"
    existing_tests: Optional[str] = None


class GeneratedTests(CapabilityResponse):
    test_cases: list[str] = Field(default_factory=list)
    fixtures_needed: list[str] = Field(default_factory=list)
    coverage_notes: str = ""


class RequirementDecomposeRequest(CapabilityRequest):
    requirement_text: str = Field(..., min_length=1)
    context: Optional[str] = None
    existing_subtasks: Sequence[str] = Field(default_factory=list)
    max_subtasks: int = Field(20, ge=1, le=100)


class SubTask(BaseModel):
    title: str
    description: str
    estimated_complexity: str = Field(..., pattern="^(xs|s|m|l|xl)$")
    depends_on: list[int] = Field(
        default_factory=list,
        description="Index (0-based) of sibling subtasks this one depends on.",
    )


class RequirementDecomposition(CapabilityResponse):
    subtasks: list[SubTask] = Field(default_factory=list)
    architecture_notes: list[str] = Field(default_factory=list)
    open_questions: list[str] = Field(default_factory=list)