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

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

from modules.shared.attributeUtils import register_model_labels, ModelMixin

# 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 Task(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 Task
register_model_labels(
    "Task",
    {"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 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[Task] = 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"}
    }
)

# AGENT AND TASK MODELS

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"}
    }
)

class AgentResponse(BaseModel, ModelMixin):
    """Data model for an agent response"""
    success: bool = Field(description="Whether the agent execution was successful")
    message: ChatMessage = Field(description="Response message from the agent")
    performance: Dict[str, Any] = Field(default_factory=dict, description="Performance metrics")
    progress: float = Field(description="Task progress (0-100)")

# Register labels for AgentResponse
register_model_labels(
    "AgentResponse",
    {"en": "Agent Response", "fr": "Réponse de l'agent"},
    {
        "success": {"en": "Success", "fr": "Succès"},
        "message": {"en": "Message", "fr": "Message"},
        "performance": {"en": "Performance", "fr": "Performance"},
        "progress": {"en": "Progress", "fr": "Progression"}
    }
)

class TaskPlan(BaseModel, ModelMixin):
    """Data model for a task plan"""
    fileList: List[str] = Field(default_factory=list, description="List of files")
    tasks: List[Task] = Field(default_factory=list, description="List of tasks in the plan")
    userLanguage: str = Field(description="User's preferred language")
    userResponse: str = Field(description="User's response or feedback")

# Register labels for TaskPlan
register_model_labels(
    "TaskPlan",
    {"en": "Task Plan", "fr": "Plan de tâches"},
    {
        "fileList": {"en": "File List", "fr": "Liste de fichiers"},
        "tasks": {"en": "Tasks", "fr": "Tâches"},
        "userLanguage": {"en": "User Language", "fr": "Langue de l'utilisateur"},
        "userResponse": {"en": "User Response", "fr": "Réponse de l'utilisateur"}
    }
)

class UserInputRequest(BaseModel, ModelMixin):
    """Data model for a user input request"""
    prompt: str = Field(description="Prompt for the user")
    listFileId: List[int] = 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"}
    }
)

class AgentProfile(BaseModel, ModelMixin):
    """Model for agent profile information."""
    id: str
    name: str
    description: str
    capabilities: List[str] = Field(default_factory=list)
    isAvailable: bool = True
    lastActive: Optional[datetime] = None
    stats: Optional[Dict[str, Any]] = None

# Register labels for AgentProfile
register_model_labels(
    "AgentProfile",
    {"en": "Agent Profile", "fr": "Profil de l'agent"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "name": {"en": "Name", "fr": "Nom"},
        "description": {"en": "Description", "fr": "Description"},
        "capabilities": {"en": "Capabilities", "fr": "Capacités"},
        "isAvailable": {"en": "Available", "fr": "Disponible"},
        "lastActive": {"en": "Last Active", "fr": "Dernière activité"},
        "stats": {"en": "Statistics", "fr": "Statistiques"}
    }
)