"""
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")
    successRate: Optional[float] = Field(None, description="Success rate of operations")
    errorCount: Optional[int] = Field(None, description="Number of errors encountered")

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")
    performance: Optional[Dict[str, Any]] = Field(None, description="Performance metrics")

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")
    success: Optional[bool] = Field(None, description="Whether the message processing was successful")

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")
    tasks: List['Task'] = Field(default_factory=list, description="List of tasks in the workflow")
    
    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"}),
        "tasks": Label(default="Tasks", translations={"en": "Tasks", "fr": "Tâches"})
    }

# 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")
    performance: Optional[Dict[str, Any]] = Field(None, description="Performance metrics")

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

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

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

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")
    userLanguage: str = Field(description="User's preferred language")