"""
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 Label, BaseModelWithUI


# WORKFLOW MODELS

class ChatContent(BaseModelWithUI):
    """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")

class ChatDocument(BaseModelWithUI):
    """Data model for a chat document"""
    id: str = Field(description="Primary key")
    fileId: int = 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")

class ChatStat(BaseModelWithUI):
    """Data model for chat statistics"""
    id: str = Field(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")

class ChatLog(BaseModelWithUI):
    """Data model for a chat log"""
    id: str = Field(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")

class ChatMessage(BaseModelWithUI):
    """Data model for a chat message"""
    id: str = Field(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")

class ChatWorkflow(BaseModelWithUI):
    """Data model for a chat workflow"""
    id: str = Field(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")
    
    label: Label = Field(
        default=Label(default="Chat Workflow", translations={"en": "Chat Workflow", "fr": "Flux de travail de chat"}),
        description="Label for the class"
    )
    
    fieldLabels: Dict[str, Label] = {
        "id": Label(default="ID", translations={}),
        "mandateId": Label(default="Mandate ID", translations={"en": "Mandate ID", "fr": "ID du mandat"}),
        "status": Label(default="Status", translations={"en": "Status", "fr": "Statut"}),
        "name": Label(default="Name", translations={"en": "Name", "fr": "Nom"}),
        "currentRound": Label(default="Current Round", translations={"en": "Current Round", "fr": "Tour actuel"}),
        "lastActivity": Label(default="Last Activity", translations={"en": "Last Activity", "fr": "Dernière activité"}),
        "startedAt": Label(default="Started At", translations={"en": "Started At", "fr": "Démarré le"}),
        "logs": Label(default="Logs", translations={"en": "Logs", "fr": "Journaux"}),
        "messages": Label(default="Messages", translations={"en": "Messages", "fr": "Messages"}),
        "stats": Label(default="Statistics", translations={"en": "Statistics", "fr": "Statistiques"})
    }

# AGENT AND TASK MODELS

class Agent(BaseModelWithUI):
    """Data model for an agent"""
    id: str = Field(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")

class AgentResponse(BaseModelWithUI):
    """Data model for an agent response"""
    response: str = Field(description="Response content from the agent")
    documents: List[ChatDocument] = Field(default_factory=list, description="Documents associated with the response")

class TaskItem(BaseModelWithUI):
    """Data model for a task item"""
    sequenceNr: int = Field(description="Sequence number of the task")
    agentName: str = Field(description="Name of the agent assigned to this task")
    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 (format: filename;[documentId])")
    filesOutput: List[str] = Field(default_factory=list, description="Output files (format: filename)")

class TaskPlan(BaseModelWithUI):
    """Data model for a task plan"""
    fileList: List[str] = Field(default_factory=list, description="List of files (format: filename)")
    taskItems: List[TaskItem] = Field(default_factory=list, description="List of task items in the plan")
    userLanguage: str = Field(description="User's preferred language")
    userResponse: str = Field(description="User's response or feedback")

class UserInputRequest(BaseModelWithUI):
    """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")