"""
Chat model classes for the chat system.
"""

from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional, Union
from datetime import datetime
import uuid

from modules.shared.attributeUtils import register_model_labels, ModelMixin


# USER MODELS

class UserInputRequest(BaseModel, ModelMixin):
    """Data model for a user input request"""
    prompt: str = Field(description="Prompt for the user")
    listFileId: List[str] = Field(default_factory=list, description="List of file IDs")
    userLanguage: str = Field(default="en", description="User's preferred language")
# Register labels for UserInputRequest
register_model_labels(
    "UserInputRequest",
    {"en": "User Input Request", "fr": "Demande de saisie utilisateur"},
    {
        "prompt": {"en": "Prompt", "fr": "Invite"},
        "listFileId": {"en": "File IDs", "fr": "IDs des fichiers"},
        "userLanguage": {"en": "User Language", "fr": "Langue de l'utilisateur"}
    }
)

# WORKFLOW MODELS

class ChatContent(BaseModel, ModelMixin):
    """Data model for chat content"""
    sequenceNr: int = Field(description="Sequence number of the content")
    name: str = Field(description="Name of the content")
    data: str = Field(description="The actual content data")
    mimeType: str = Field(description="MIME type of the content")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
# Register labels for ChatContent
register_model_labels(
    "ChatContent",
    {"en": "Chat Content", "fr": "Contenu de chat"},
    {
        "sequenceNr": {"en": "Sequence Number", "fr": "Numéro de séquence"},
        "name": {"en": "Name", "fr": "Nom"},
        "data": {"en": "Data", "fr": "Données"},
        "mimeType": {"en": "MIME Type", "fr": "Type MIME"},
        "metadata": {"en": "Metadata", "fr": "Métadonnées"}
    }
)

class ChatDocument(BaseModel, ModelMixin):
    """Data model for a chat document"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key")
    fileId: str = Field(description="Foreign key to file")
    filename: str = Field(description="Name of the file")
    fileSize: int = Field(description="Size of the file")
    mimeType: str = Field(description="MIME type of the file")
    contents: List[ChatContent] = Field(default_factory=list, description="List of chat contents")
# Register labels for ChatDocument
register_model_labels(
    "ChatDocument",
    {"en": "Chat Document", "fr": "Document de chat"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "fileId": {"en": "File ID", "fr": "ID du fichier"},
        "filename": {"en": "Filename", "fr": "Nom de fichier"},
        "fileSize": {"en": "File Size", "fr": "Taille du fichier"},
        "mimeType": {"en": "MIME Type", "fr": "Type MIME"},
        "contents": {"en": "Contents", "fr": "Contenus"}
    }
)

class ChatStat(BaseModel, ModelMixin):
    """Data model for chat statistics"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key")
    processingTime: Optional[float] = Field(None, description="Processing time in seconds")
    tokenCount: Optional[int] = Field(None, description="Number of tokens processed")
    bytesSent: Optional[int] = Field(None, description="Number of bytes sent")
    bytesReceived: Optional[int] = Field(None, description="Number of bytes received")
    successRate: Optional[float] = Field(None, description="Success rate of operations")
    errorCount: Optional[int] = Field(None, description="Number of errors encountered")
# Register labels for ChatStat
register_model_labels(
    "ChatStat",
    {"en": "Chat Statistics", "fr": "Statistiques de chat"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "processingTime": {"en": "Processing Time", "fr": "Temps de traitement"},
        "tokenCount": {"en": "Token Count", "fr": "Nombre de tokens"},
        "bytesSent": {"en": "Bytes Sent", "fr": "Octets envoyés"},
        "bytesReceived": {"en": "Bytes Received", "fr": "Octets reçus"},
        "successRate": {"en": "Success Rate", "fr": "Taux de succès"},
        "errorCount": {"en": "Error Count", "fr": "Nombre d'erreurs"}
    }
)

class ChatLog(BaseModel, ModelMixin):
    """Data model for a chat log"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key")
    workflowId: str = Field(description="Foreign key to workflow")
    message: str = Field(description="Log message")
    type: str = Field(description="Type of log entry")
    timestamp: str = Field(description="Timestamp of the log entry")
    agentName: str = Field(description="Name of the agent")
    status: str = Field(description="Status of the log entry")
    progress: Optional[int] = Field(None, description="Progress percentage")
    performance: Optional[Dict[str, Any]] = Field(None, description="Performance metrics")
# Register labels for ChatLog
register_model_labels(
    "ChatLog",
    {"en": "Chat Log", "fr": "Journal de chat"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "workflowId": {"en": "Workflow ID", "fr": "ID du flux de travail"},
        "message": {"en": "Message", "fr": "Message"},
        "type": {"en": "Type", "fr": "Type"},
        "timestamp": {"en": "Timestamp", "fr": "Horodatage"},
        "agentName": {"en": "Agent Name", "fr": "Nom de l'agent"},
        "status": {"en": "Status", "fr": "Statut"},
        "progress": {"en": "Progress", "fr": "Progression"},
        "performance": {"en": "Performance", "fr": "Performance"}
    }
)

class ChatMessage(BaseModel, ModelMixin):
    """Data model for a chat message"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key")
    workflowId: str = Field(description="Foreign key to workflow")
    parentMessageId: Optional[str] = Field(None, description="Parent message ID for threading")
    agentName: Optional[str] = Field(None, description="Name of the agent")
    documents: List[ChatDocument] = Field(default_factory=list, description="Associated documents")
    message: Optional[str] = Field(None, description="Message content")
    role: str = Field(description="Role of the message sender")
    status: str = Field(description="Status of the message")
    sequenceNr: int = Field(description="Sequence number of the message")
    startedAt: str = Field(description="When the message processing started")
    finishedAt: Optional[str] = Field(None, description="When the message processing finished")
    stats: Optional[ChatStat] = Field(None, description="Statistics for this message")
    success: Optional[bool] = Field(None, description="Whether the message processing was successful")
# Register labels for ChatMessage
register_model_labels(
    "ChatMessage",
    {"en": "Chat Message", "fr": "Message de chat"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "workflowId": {"en": "Workflow ID", "fr": "ID du flux de travail"},
        "parentMessageId": {"en": "Parent Message ID", "fr": "ID du message parent"},
        "agentName": {"en": "Agent Name", "fr": "Nom de l'agent"},
        "documents": {"en": "Documents", "fr": "Documents"},
        "message": {"en": "Message", "fr": "Message"},
        "role": {"en": "Role", "fr": "Rôle"},
        "status": {"en": "Status", "fr": "Statut"},
        "sequenceNr": {"en": "Sequence Number", "fr": "Numéro de séquence"},
        "startedAt": {"en": "Started At", "fr": "Démarré le"},
        "finishedAt": {"en": "Finished At", "fr": "Terminé le"},
        "stats": {"en": "Statistics", "fr": "Statistiques"},
        "success": {"en": "Success", "fr": "Succès"}
    }
)

class AgentTask(BaseModel, ModelMixin):
    """Data model for a task"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key")
    workflowId: str = Field(description="Foreign key to workflow")
    agentName: str = Field(description="Name of the agent assigned to this task")
    status: str = Field(description="Current status of the task")
    progress: float = Field(description="Task progress (0-100)")
    prompt: str = Field(description="Prompt for the task")
    userLanguage: str = Field(description="User's preferred language")
    filesInput: List[str] = Field(default_factory=list, description="Input files")
    filesOutput: List[str] = Field(default_factory=list, description="Output files")
    result: Optional[ChatMessage] = Field(None, description="Task result message")
    error: Optional[str] = Field(None, description="Error message if failed")
    startedAt: str = Field(description="When the task started")
    finishedAt: Optional[str] = Field(None, description="When the task finished")
    performance: Optional[Dict[str, Any]] = Field(None, description="Performance metrics")
# Register labels for AgentTask
register_model_labels(
    "AgentTask",
    {"en": "Task", "fr": "Tâche"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "workflowId": {"en": "Workflow ID", "fr": "ID du flux de travail"},
        "agentName": {"en": "Agent Name", "fr": "Nom de l'agent"},
        "status": {"en": "Status", "fr": "Statut"},
        "progress": {"en": "Progress", "fr": "Progression"},
        "prompt": {"en": "Prompt", "fr": "Invite"},
        "userLanguage": {"en": "User Language", "fr": "Langue de l'utilisateur"},
        "filesInput": {"en": "Input Files", "fr": "Fichiers d'entrée"},
        "filesOutput": {"en": "Output Files", "fr": "Fichiers de sortie"},
        "result": {"en": "Result", "fr": "Résultat"},
        "error": {"en": "Error", "fr": "Erreur"},
        "startedAt": {"en": "Started At", "fr": "Démarré le"},
        "finishedAt": {"en": "Finished At", "fr": "Terminé le"},
        "performance": {"en": "Performance", "fr": "Performance"}
    }
)

class Agent(BaseModel, ModelMixin):
    """Data model for an agent"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key")
    name: str = Field(description="Name of the agent")
    description: str = Field(description="Description of the agent")
    capabilities: List[str] = Field(default_factory=list, description="List of agent capabilities")
    performance: Optional[Dict[str, Any]] = Field(None, description="Performance metrics")
# Register labels for Agent
register_model_labels(
    "Agent",
    {"en": "Agent", "fr": "Agent"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "name": {"en": "Name", "fr": "Nom"},
        "description": {"en": "Description", "fr": "Description"},
        "capabilities": {"en": "Capabilities", "fr": "Capacités"},
        "performance": {"en": "Performance", "fr": "Performance"}
    }
)

# WORKFLOW MODELS

class ChatWorkflow(BaseModel, ModelMixin):
    """Data model for a chat workflow"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key")
    mandateId: str = Field(description="ID of the mandate this workflow belongs to")
    status: str = Field(description="Current status of the workflow")
    name: Optional[str] = Field(None, description="Name of the workflow")
    currentRound: int = Field(description="Current round number")
    lastActivity: str = Field(description="Timestamp of last activity")
    startedAt: str = Field(description="When the workflow started")
    logs: List[ChatLog] = Field(default_factory=list, description="Workflow logs")
    messages: List[ChatMessage] = Field(default_factory=list, description="Messages in the workflow")
    stats: Optional[ChatStat] = Field(None, description="Workflow statistics")
    tasks: List[AgentTask] = Field(default_factory=list, description="List of tasks in the workflow")
# Register labels for ChatWorkflow
register_model_labels(
    "ChatWorkflow",
    {"en": "Chat Workflow", "fr": "Flux de travail de chat"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "mandateId": {"en": "Mandate ID", "fr": "ID du mandat"},
        "status": {"en": "Status", "fr": "Statut"},
        "name": {"en": "Name", "fr": "Nom"},
        "currentRound": {"en": "Current Round", "fr": "Tour actuel"},
        "lastActivity": {"en": "Last Activity", "fr": "Dernière activité"},
        "startedAt": {"en": "Started At", "fr": "Démarré le"},
        "logs": {"en": "Logs", "fr": "Journaux"},
        "messages": {"en": "Messages", "fr": "Messages"},
        "stats": {"en": "Statistics", "fr": "Statistiques"},
        "tasks": {"en": "Tasks", "fr": "Tâches"}
    }
)

# DOCUMENT MODELS

class DocumentExtraction(BaseModel, ModelMixin):
    """Data model for document extraction history"""
    timestamp: str = Field(description="Timestamp of extraction")
    type: str = Field(description="Type of document")
    sections: List[str] = Field(default_factory=list, description="Extracted sections")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Extraction metadata")

# Register labels for DocumentExtraction
register_model_labels(
    "DocumentExtraction",
    {"en": "Document Extraction", "fr": "Extraction de document"},
    {
        "timestamp": {"en": "Timestamp", "fr": "Horodatage"},
        "type": {"en": "Type", "fr": "Type"},
        "sections": {"en": "Sections", "fr": "Sections"},
        "metadata": {"en": "Metadata", "fr": "Métadonnées"}
    }
)

class DocumentContext(BaseModel, ModelMixin):
    """Data model for document context"""
    id: str = Field(description="Document ID")
    extractionHistory: List[DocumentExtraction] = Field(default_factory=list, description="History of extractions")
    relevantSections: List[str] = Field(default_factory=list, description="Relevant sections")
    processingStatus: Dict[str, str] = Field(default_factory=dict, description="Processing status")

# Register labels for DocumentContext
register_model_labels(
    "DocumentContext",
    {"en": "Document Context", "fr": "Contexte de document"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "extractionHistory": {"en": "Extraction History", "fr": "Historique d'extraction"},
        "relevantSections": {"en": "Relevant Sections", "fr": "Sections pertinentes"},
        "processingStatus": {"en": "Processing Status", "fr": "Statut de traitement"}
    }
)

class DocumentMetadata(BaseModel, ModelMixin):
    """Data model for document metadata"""
    type: str = Field(description="Document type")
    format: str = Field(description="Document format")
    size: int = Field(description="Document size in bytes")
    pages: Optional[int] = Field(None, description="Number of pages")
    sections: Optional[List[str]] = Field(None, description="Document sections")
    error: Optional[str] = Field(None, description="Processing error if any")

# Register labels for DocumentMetadata
register_model_labels(
    "DocumentMetadata",
    {"en": "Document Metadata", "fr": "Métadonnées de document"},
    {
        "type": {"en": "Type", "fr": "Type"},
        "format": {"en": "Format", "fr": "Format"},
        "size": {"en": "Size", "fr": "Taille"},
        "pages": {"en": "Pages", "fr": "Pages"},
        "sections": {"en": "Sections", "fr": "Sections"},
        "error": {"en": "Error", "fr": "Erreur"}
    }
)

class ImageData(BaseModel, ModelMixin):
    """Data model for image data"""
    data: str = Field(description="Base64 encoded image data")
    format: str = Field(description="Image format")
    page: Optional[int] = Field(None, description="Page number if from a multi-page document")
    index: Optional[int] = Field(None, description="Image index in the document")

# Register labels for ImageData
register_model_labels(
    "ImageData",
    {"en": "Image Data", "fr": "Données d'image"},
    {
        "data": {"en": "Image Data", "fr": "Données d'image"},
        "format": {"en": "Format", "fr": "Format"},
        "page": {"en": "Page", "fr": "Page"},
        "index": {"en": "Index", "fr": "Index"}
    }
)

class DocumentContent(BaseModel, ModelMixin):
    """Data model for document content"""
    text: Optional[str] = Field(None, description="Extracted text content")
    data: Optional[Dict[str, Any]] = Field(None, description="Structured data content")
    images: Optional[List[ImageData]] = Field(None, description="Extracted images")
    metadata: DocumentMetadata = Field(description="Document metadata")

# Register labels for DocumentContent
register_model_labels(
    "DocumentContent",
    {"en": "Document Content", "fr": "Contenu de document"},
    {
        "text": {"en": "Text", "fr": "Texte"},
        "data": {"en": "Data", "fr": "Données"},
        "images": {"en": "Images", "fr": "Images"},
        "metadata": {"en": "Metadata", "fr": "Métadonnées"}
    }
)

class ProcessedDocument(BaseModel, ModelMixin):
    """Data model for processed document"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Document ID")
    name: str = Field(description="Document name")
    contentType: str = Field(description="Content type")
    content: DocumentContent = Field(description="Document content")
    context: Optional[DocumentContext] = Field(None, description="Document context")

# Register labels for ProcessedDocument
register_model_labels(
    "ProcessedDocument",
    {"en": "Processed Document", "fr": "Document traité"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "name": {"en": "Name", "fr": "Nom"},
        "contentType": {"en": "Content Type", "fr": "Type de contenu"},
        "content": {"en": "Content", "fr": "Contenu"},
        "context": {"en": "Context", "fr": "Contexte"}
    }
)