#!/usr/bin/env python3
"""
Structured outputs for Agent API
"""

from typing import List, Optional
from pydantic import BaseModel, Field


class Message(BaseModel):
    """Chat message"""
    role: str = Field(..., description="Role: 'user' or 'assistant'")
    content: str = Field(..., description="Message content")


class LawyerProfile(BaseModel):
    """Lawyer's professional profile"""
    full_name: Optional[str] = Field(None, description="Lawyer's full name")
    primary_specialty: Optional[str] = Field(None, description="Primary legal specialty")
    legal_specialties: Optional[List[str]] = Field(None, description="List of legal specialties")
    jurisdiction: Optional[str] = Field(None, description="Primary jurisdiction of practice")
    experience_level: Optional[str] = Field(None, description="Experience level (e.g., Junior, Mid-level, Senior, Partner)")
    languages: Optional[List[str]] = Field(None, description="Languages spoken")
    lawyer_description: Optional[str] = Field(None, description="Professional description or bio")


class DocumentAnalysis(BaseModel):
    """Document analysis result"""
    file_name: str
    summary: Optional[str] = None
    actors: Optional[str] = None
    key_details: Optional[str] = None


class ChatRequest(BaseModel):
    """Chat request model"""
    clientId: str = Field(..., description="Unique client identifier")
    message: str = Field(..., description="User's question")
    conversationHistory: Optional[List[Message]] = Field(default=[], description="Previous conversation messages")
    userType: Optional[str] = Field(default="client", description="User type: 'client' for general users or 'lawyer' for legal professionals")
    jurisdiction: Optional[str] = Field(default="Romania", description="Jurisdiction of the user")
    documentAnalyses: Optional[List[DocumentAnalysis]] = Field(default=None, description="Lawyer's document analyses")
    lawyerProfile: Optional[LawyerProfile] = Field(default=None, description="Lawyer's professional profile")


class ChatResponse(BaseModel):
    """Chat response model"""
    response: str = Field(..., description="Assistant's response")
    processing_time: float = Field(..., description="Processing time in seconds")
    references: List[str] = Field(default=[], description="Referenced documents")
    timestamp: str = Field(..., description="Response timestamp")
    error: Optional[str] = Field(None, description="Error message if any")


class HealthResponse(BaseModel):
    """Health check response model"""
    status: str = Field(..., description="Health status")
    agent_ready: bool = Field(..., description="Whether agent is ready")
    lightrag_healthy: bool = Field(..., description="Whether LightRAG is healthy")
    timestamp: str = Field(..., description="Health check timestamp")


class AnalyzePDFRequest(BaseModel):
    """PDF analysis request model"""
    pdf_content: str = Field(..., description="Base64 encoded document content (PDF or image)")
    filename: Optional[str] = Field(default="document.pdf", description="Original filename")


class AnalyzePDFResponse(BaseModel):
    """PDF analysis response model"""
    actors: str = Field(..., description="Extracted actors")
    key_details: str = Field(..., description="Key details extracted")
    summary: str = Field(..., description="High-level summary")
    processing_status: str = Field(..., description="Processing status")
    processing_time: float = Field(..., description="Processing time in seconds")
    timestamp: str = Field(..., description="Analysis timestamp")
    error: Optional[str] = Field(None, description="Error message if any")