schemas.py

from __future__ import annotations

import uuid
from datetime import datetime
from typing import Optional

from pydantic import BaseModel, Field

from backend.models.enums import AgentPhase, AmbiguityLevel, ContextSource, PipelineStatus, WSMessageType


# ---------------------------------------------------------------------------
# Agent I/O schemas (inter-agent message protocol)
# ---------------------------------------------------------------------------

class IntentSchema(BaseModel):
    """Output of Phase 1 / Input to Phase 2."""
    goal: str
    context: str = ""
    constraints: list[str] = Field(default_factory=list)
    expected_output_format: str = ""
    success_criteria: list[str] = Field(default_factory=list)
    alternative_interpretations: list[str] = Field(default_factory=list)
    confidence: float = Field(ge=0.0, le=1.0, default=0.5)
    raw_input: str = ""


class AmbiguityFlag(BaseModel):
    dimension: str
    description: str
    level: AmbiguityLevel = AmbiguityLevel.LOW
    impact: str = ""


class AmbiguityReport(BaseModel):
    """Output of Phase 2 / Input to Phase 3."""
    intent: IntentSchema
    flags: list[AmbiguityFlag] = Field(default_factory=list)
    overall_ambiguity: AmbiguityLevel = AmbiguityLevel.LOW
    needs_clarification: bool = False
    reasoning: str = ""


class ClarificationQuestion(BaseModel):
    question: str
    why_needed: str = ""
    dimension: str = ""
    priority: int = 1
    options: list[str] = Field(default_factory=list)
    auto_resolved: bool = False
    auto_answer: str = ""
    auto_answer_source: str = ""


class ClarificationRequest(BaseModel):
    """Output of Phase 3 when questions are needed."""
    questions: list[ClarificationQuestion] = Field(default_factory=list)
    search_context: list[dict] = Field(
        default_factory=list,
        description="Web search results that informed the questions (title, url, snippet).",
    )
    auto_resolved_count: int = 0


class ClarifiedIntent(BaseModel):
    """Output of Phase 3 after user answers / Input to Phase 4."""
    original_intent: IntentSchema
    clarifications: dict[str, str] = Field(default_factory=dict)
    refined_goal: str = ""
    refined_constraints: list[str] = Field(default_factory=list)
    refined_output_format: str = ""


class PlanStep(BaseModel):
    step_number: int
    description: str
    inputs: list[str] = Field(default_factory=list)
    expected_output: str = ""
    dependencies: list[int] = Field(default_factory=list)
    validation: str = ""


class ExecutionPlan(BaseModel):
    """Output of Phase 4 / Input to Phase 5."""
    objective: str = ""
    inputs: list[str] = Field(default_factory=list)
    steps: list[PlanStep] = Field(default_factory=list)
    deliverables: list[str] = Field(default_factory=list)
    validation_checkpoints: list[str] = Field(default_factory=list)


class ContextItem(BaseModel):
    """A single retrieved context entry from one source."""
    source: ContextSource
    query: str = ""
    content: str = ""
    relevance: str = ""
    source_detail: str = ""       # doc name, URL, phase name, or step number
    reasoning_trace: str = ""     # PageIndex hierarchical reasoning path
    metadata: dict = Field(default_factory=dict)


class ToolRequirement(BaseModel):
    """Describes one retrieval tool needed to gather context for a step."""
    tool_name: str                # "pageindex", "exa_search", "prior_context", "step_dep"
    purpose: str = ""
    queries: list[str] = Field(default_factory=list)
    doc_filter: list[str] = Field(default_factory=list)   # PageIndex: which docs
    priority: str = "required"    # "required" | "optional"


class StepContext(BaseModel):
    """All retrieved context for a single execution step."""
    step_number: int
    step_description: str = ""
    required_tools: list[ToolRequirement] = Field(default_factory=list)
    gathered_contexts: list[ContextItem] = Field(default_factory=list)
    context_summary: str = ""     # LLM-synthesised briefing for this step
    execution_strategy: str = ""  # recommended approach given gathered context


class ResourcePlan(BaseModel):
    """Structured resource and context plan for the full execution."""
    objective: str = ""
    total_steps: int = 0
    tool_inventory: list[dict] = Field(default_factory=list)
    step_contexts: list[StepContext] = Field(default_factory=list)
    dependency_graph: dict = Field(default_factory=dict)
    knowledge_base_docs_used: list[str] = Field(default_factory=list)
    web_searches_performed: int = 0
    estimated_complexity: str = ""


class StepResult(BaseModel):
    step_number: int
    output: str = ""
    status: str = "completed"
    validation_passed: bool = True
    contexts_used: list[str] = Field(default_factory=list)
    reasoning: str = ""


class ExecutionResult(BaseModel):
    """Output of Phase 5 — the final deliverable."""
    plan: ExecutionPlan | None = None
    resource_plan: ResourcePlan | None = None
    step_results: list[StepResult] = Field(default_factory=list)
    final_output: str = ""
    completeness_check: str = ""
    clarity_check: str = ""
    relevance_check: str = ""
    correctness_check: str = ""
    trace_to_goal: str = ""
    blob_url: str = ""  # Vercel Blob public URL of the rendered HTML output


# ---------------------------------------------------------------------------
# WebSocket message envelope
# ---------------------------------------------------------------------------

class WSMessage(BaseModel):
    type: WSMessageType
    session_id: str = ""
    phase: Optional[AgentPhase] = None
    data: dict = Field(default_factory=dict)
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    message_id: str = Field(default_factory=lambda: str(uuid.uuid4()))


# ---------------------------------------------------------------------------
# Session state
# ---------------------------------------------------------------------------

class SessionState(BaseModel):
    session_id: str
    status: PipelineStatus = PipelineStatus.RUNNING
    current_phase: Optional[AgentPhase] = None
    user_input: str = ""
    intent: Optional[IntentSchema] = None
    ambiguity_report: Optional[AmbiguityReport] = None
    clarification_request: Optional[ClarificationRequest] = None
    clarified_intent: Optional[ClarifiedIntent] = None
    execution_plan: Optional[ExecutionPlan] = None
    resource_plan: Optional[ResourcePlan] = None
    execution_result: Optional[ExecutionResult] = None
    created_at: datetime = Field(default_factory=datetime.utcnow)


# ---------------------------------------------------------------------------
# Artifact
# ---------------------------------------------------------------------------

class Artifact(BaseModel):
    artifact_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    session_id: str
    artifact_type: str
    title: str
    content: dict = Field(default_factory=dict)
    created_at: datetime = Field(default_factory=datetime.utcnow)