resolutionflow/backend/app/schemas/ai_session.py

"""Pydantic schemas for FlowPilot AI sessions."""
from __future__ import annotations

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

from pydantic import BaseModel, Field


# ── Intake ──

class AISessionCreateRequest(BaseModel):
    """Start a new FlowPilot or chat session."""
    session_type: str = Field(
        "guided",
        pattern="^(guided|chat)$",
        description="Session type: guided (FlowPilot) or chat (assistant)",
    )
    intake_type: str = Field(
        "free_text",
        pattern="^(free_text|psa_ticket|screenshot|log_paste|combined)$",
    )
    intake_content: dict[str, Any] = Field(
        ...,
        description=(
            "Intake payload. Shape depends on intake_type: "
            "{text: str} for free_text, "
            "{text?: str, image_urls?: list[str]} for screenshot, "
            "{text?: str, log_content?: str} for log_paste, "
            "{ticket_id: str, psa_connection_id: str} for psa_ticket, "
            "any combination for combined."
        ),
    )
    psa_ticket_id: Optional[str] = None
    psa_connection_id: Optional[UUID] = None


class AISessionCreateResponse(BaseModel):
    """Response after starting a session — includes the first FlowPilot step."""
    session_id: UUID
    status: str
    confidence_tier: str
    problem_summary: str | None = None
    problem_domain: str | None = None
    matched_flow_id: UUID | None = None
    matched_flow_name: str | None = None
    match_score: float | None = None
    first_step: AISessionStepResponse
    psa_context_status: str | None = None  # loaded | unavailable | None (no PSA)


# ── Step interaction ──

class StepOptionSchema(BaseModel):
    """A selectable option presented to the engineer."""
    label: str
    value: str
    followup_hint: str | None = None


class AISessionStepResponse(BaseModel):
    """A FlowPilot step rendered in the session UI."""
    step_id: UUID
    step_order: int
    step_type: str
    content: dict[str, Any]
    context_message: str | None = None
    options: list[StepOptionSchema] = []
    allow_free_text: bool = True
    allow_skip: bool = True
    confidence_tier: str
    confidence_score: float

    model_config = {"from_attributes": True}


class StepResponseRequest(BaseModel):
    """Engineer's response to a FlowPilot step."""
    selected_option: str | None = None
    free_text_input: str | None = None
    was_skipped: bool = False
    action_result: dict[str, Any] | None = None


class StepResponseResponse(BaseModel):
    """FlowPilot's next step after processing the engineer's response."""
    session_id: UUID
    status: str
    confidence_tier: str
    confidence_score: float
    next_step: AISessionStepResponse | None = None
    resolution_suggested: bool = False
    resolution_summary: str | None = None


# ── Resolution / Escalation ──

class ResolveSessionRequest(BaseModel):
    """Close a session as resolved."""
    resolution_summary: str = Field(..., min_length=5, max_length=2000)
    resolution_action: str | None = None
    session_rating: int | None = Field(None, ge=1, le=5)
    session_feedback: str | None = None


class EscalateSessionRequest(BaseModel):
    """Escalate a session to another engineer."""
    escalation_reason: str = Field(..., min_length=5, max_length=2000)
    escalated_to_id: UUID | None = None


class DocumentationStep(BaseModel):
    """A step in the documentation trail."""
    step_number: int
    step_type: str
    description: str
    engineer_response: str | None = None
    outcome: str | None = None


class SessionDocumentation(BaseModel):
    """Auto-generated session documentation."""
    problem_summary: str
    problem_domain: str | None = None
    intake_summary: str
    diagnostic_steps: list[DocumentationStep]
    resolution_summary: str | None = None
    escalation_reason: str | None = None
    total_steps: int
    duration_display: str | None = None
    generated_at: datetime


class SessionCloseResponse(BaseModel):
    """Response after resolving or escalating."""
    session_id: UUID
    status: str
    documentation: SessionDocumentation
    psa_push_status: str = "no_psa"  # sent | pending_retry | no_psa | failed
    psa_push_error: str | None = None
    member_mapping_warning: str | None = None


class StatusUpdateRequest(BaseModel):
    """Generate a mid-session or post-session status update."""
    audience: str = Field(
        ...,
        pattern="^(ticket_notes|client_update|email_draft)$",
        description="Who is this update for?",
    )
    length: str = Field(
        "detailed",
        pattern="^(quick|detailed)$",
        description="Quick (1-2 sentences) or detailed breakdown",
    )
    context: str = Field(
        "status",
        pattern="^(status|resolution|escalation)$",
        description="What type of communication: mid-session status, resolution close-out, or escalation handoff",
    )


class StatusUpdateResponse(BaseModel):
    """Generated status update content."""
    content: str
    audience: str
    length: str
    context: str
    session_status: str
    steps_completed: int
    time_spent_display: str | None = None
    client_name: str | None = None
    generated_at: datetime


class RateSessionRequest(BaseModel):
    """Submit post-session rating."""
    rating: int = Field(..., ge=1, le=5)
    feedback: str | None = None


class PickupSessionRequest(BaseModel):
    """Pick up an escalated session as a new engineer."""
    resume_mode: str = Field("continue", pattern="^(continue|fresh)$")
    additional_context: str | None = None


class LinkTicketRequest(BaseModel):
    """Link a PSA ticket to an in-progress session."""
    psa_ticket_id: str
    psa_connection_id: UUID


# ── List / Detail ──

class AISessionSummary(BaseModel):
    """Compact session for list views."""
    id: UUID
    session_type: str = "guided"
    title: str | None = None
    status: str
    intake_type: str
    problem_summary: str | None = None
    problem_domain: str | None = None
    confidence_tier: str
    step_count: int
    session_rating: int | None = None
    psa_ticket_id: str | None = None
    escalation_reason: str | None = None
    created_at: datetime
    resolved_at: datetime | None = None

    model_config = {"from_attributes": True}


class AISessionDetail(AISessionSummary):
    """Full session detail with steps (guided) or messages (chat)."""
    intake_content: dict[str, Any]
    matched_flow_id: UUID | None = None
    match_score: float | None = None
    resolution_summary: str | None = None
    resolution_action: str | None = None
    escalation_reason: str | None = None
    session_feedback: str | None = None
    psa_ticket_id: str | None = None
    psa_connection_id: UUID | None = None
    ticket_data: dict[str, Any] | None = None
    steps: list[AISessionStepResponse] = []
    conversation_messages: list[dict[str, Any]] = []  # Chat sessions store messages here
    pending_task_lane: dict[str, Any] | None = None
    is_branching: bool = False
    active_branch_id: str | None = None

    model_config = {"from_attributes": True}


# ── Chat session ──

class ChatSessionCreateResponse(BaseModel):
    """Response after creating a chat session on ai_sessions."""
    session_id: UUID
    session_type: str = "chat"
    title: str
    status: str = "active"


class ChatMessageRequest(BaseModel):
    """Send a message in a chat session."""
    message: str = Field(..., min_length=1, max_length=8000)
    upload_ids: list[UUID] = Field(default_factory=list, max_length=10)


class ForkBranchInfo(BaseModel):
    """Branch info returned when a fork is created."""
    branch_id: str
    label: str


class ForkMetadata(BaseModel):
    """Metadata returned when the AI suggests a diagnostic fork."""
    fork_point_id: str
    fork_reason: str
    branches: list[ForkBranchInfo]
    active_branch_id: str


class ActionItem(BaseModel):
    """A single action item for the engineer."""
    label: str
    command: str | None = None
    description: str = ""


class QuestionItem(BaseModel):
    """A question the AI needs answered by the engineer."""
    text: str
    context: str = ""


class ChatMessageResponse(BaseModel):
    """AI response to a chat message."""
    content: str
    suggested_flows: list[dict[str, Any]] = []
    fork: ForkMetadata | None = None
    actions: list[ActionItem] | None = None
    questions: list[QuestionItem] | None = None


class SaveTaskLaneRequest(BaseModel):
    """Save the full task lane state (AI items + user responses)."""
    questions: list[QuestionItem] = []
    actions: list[ActionItem] = []
    responses: list[dict[str, Any]] = Field(
        default_factory=list,
        description="User's in-progress task responses with state/value",
    )


class AISessionSearchResult(BaseModel):
    """Lightweight session result for Command Palette / autocomplete."""
    id: UUID
    problem_summary: str | None = None
    problem_domain: str | None = None
    status: str
    created_at: datetime

    model_config = {"from_attributes": True}