# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""Chat models: ChatWorkflow, ChatMessage, ChatLog, ChatStat, ChatDocument."""

from typing import List, Dict, Any, Optional
from enum import Enum
from pydantic import BaseModel, Field
from modules.shared.attributeUtils import registerModelLabels
from modules.shared.timeUtils import getUtcTimestamp
import uuid


class ChatStat(BaseModel):
    """Statistics for chat operations. User-owned, no mandate context."""
    model_config = {"populate_by_name": True, "extra": "allow"}  # Allow DB system fields
    
    id: str = Field(
        default_factory=lambda: str(uuid.uuid4()), description="Primary key"
    )
    workflowId: Optional[str] = Field(
        None, description="Foreign key to workflow (for workflow stats)"
    )
    processingTime: Optional[float] = Field(
        None, description="Processing time in seconds"
    )
    bytesSent: Optional[int] = Field(None, description="Number of bytes sent")
    bytesReceived: Optional[int] = Field(None, description="Number of bytes received")
    errorCount: Optional[int] = Field(None, description="Number of errors encountered")
    process: Optional[str] = Field(None, description="The process that delivers the stats data (e.g. 'action.outlook.readMails', 'ai.process.document.name')")
    engine: Optional[str] = Field(None, description="The engine used (e.g. 'ai.anthropic.35', 'ai.tavily.basic', 'renderer.docx')")
    priceCHF: Optional[float] = Field(None, description="Calculated price in USD for the operation")


registerModelLabels(
    "ChatStat",
    {"en": "Chat Statistics", "fr": "Statistiques de chat"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "workflowId": {"en": "Workflow ID", "fr": "ID du workflow"},
        "processingTime": {"en": "Processing Time", "fr": "Temps de traitement"},
        "bytesSent": {"en": "Bytes Sent", "fr": "Octets envoyés"},
        "bytesReceived": {"en": "Bytes Received", "fr": "Octets reçus"},
        "errorCount": {"en": "Error Count", "fr": "Nombre d'erreurs"},
        "process": {"en": "Process", "fr": "Processus"},
        "engine": {"en": "Engine", "fr": "Moteur"},
        "priceCHF": {"en": "Price CHF", "fr": "Prix CHF"},
    },
)


class ChatLog(BaseModel):
    """Log entries for chat workflows. User-owned, no mandate context."""
    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="Log type (info, warning, error, etc.)")
    timestamp: float = Field(
        default_factory=getUtcTimestamp,
        description="When the log entry was created (UTC timestamp in seconds)",
    )
    status: Optional[str] = Field(None, description="Status of the log entry")
    progress: Optional[float] = Field(
        None, description="Progress indicator (0.0 to 1.0)"
    )
    performance: Optional[Dict[str, Any]] = Field(
        None, description="Performance metrics"
    )
    parentId: Optional[str] = Field(
        None, description="Parent operation ID (operationId of parent operation) for hierarchical display"
    )
    operationId: Optional[str] = Field(
        None, description="Operation ID to group related log entries"
    )
    roundNumber: Optional[int] = Field(None, description="Round number in workflow")
    taskNumber: Optional[int] = Field(None, description="Task number within round")
    actionNumber: Optional[int] = Field(None, description="Action number within task")


registerModelLabels(
    "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"},
        "status": {"en": "Status", "fr": "Statut"},
        "progress": {"en": "Progress", "fr": "Progression"},
        "performance": {"en": "Performance", "fr": "Performance"},
    },
)


class ChatDocument(BaseModel):
    """Documents attached to chat messages. User-owned, no mandate context."""
    id: str = Field(
        default_factory=lambda: str(uuid.uuid4()), description="Primary key"
    )
    messageId: str = Field(description="Foreign key to message")
    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")
    roundNumber: Optional[int] = Field(None, description="Round number in workflow")
    taskNumber: Optional[int] = Field(None, description="Task number within round")
    actionNumber: Optional[int] = Field(None, description="Action number within task")
    actionId: Optional[str] = Field(
        None, description="ID of the action that created this document"
    )


registerModelLabels(
    "ChatDocument",
    {"en": "Chat Document", "fr": "Document de chat"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "messageId": {"en": "Message ID", "fr": "ID du message"},
        "fileId": {"en": "File ID", "fr": "ID du fichier"},
        "fileName": {"en": "File Name", "fr": "Nom du fichier"},
        "fileSize": {"en": "File Size", "fr": "Taille du fichier"},
        "mimeType": {"en": "MIME Type", "fr": "Type MIME"},
        "roundNumber": {"en": "Round Number", "fr": "Numéro de tour"},
        "taskNumber": {"en": "Task Number", "fr": "Numéro de tâche"},
        "actionNumber": {"en": "Action Number", "fr": "Numéro d'action"},
        "actionId": {"en": "Action ID", "fr": "ID de l'action"},
    },
)


class ContentMetadata(BaseModel):
    size: int = Field(description="Content size in bytes")
    pages: Optional[int] = Field(
        None, description="Number of pages for multi-page content"
    )
    error: Optional[str] = Field(None, description="Processing error if any")
    width: Optional[int] = Field(None, description="Width in pixels for images/videos")
    height: Optional[int] = Field(
        None, description="Height in pixels for images/videos"
    )
    colorMode: Optional[str] = Field(None, description="Color mode")
    fps: Optional[float] = Field(None, description="Frames per second for videos")
    durationSec: Optional[float] = Field(
        None, description="Duration in seconds for media"
    )
    mimeType: str = Field(description="MIME type of the content")
    base64Encoded: bool = Field(description="Whether the data is base64 encoded")


registerModelLabels(
    "ContentMetadata",
    {"en": "Content Metadata", "fr": "Métadonnées du contenu"},
    {
        "size": {"en": "Size", "fr": "Taille"},
        "pages": {"en": "Pages", "fr": "Pages"},
        "error": {"en": "Error", "fr": "Erreur"},
        "width": {"en": "Width", "fr": "Largeur"},
        "height": {"en": "Height", "fr": "Hauteur"},
        "colorMode": {"en": "Color Mode", "fr": "Mode de couleur"},
        "fps": {"en": "FPS", "fr": "IPS"},
        "durationSec": {"en": "Duration", "fr": "Durée"},
        "mimeType": {"en": "MIME Type", "fr": "Type MIME"},
        "base64Encoded": {"en": "Base64 Encoded", "fr": "Encodé en Base64"},
    },
)


class ContentItem(BaseModel):
    label: str = Field(description="Content label")
    data: str = Field(description="Extracted text content")
    metadata: ContentMetadata = Field(description="Content metadata")


registerModelLabels(
    "ContentItem",
    {"en": "Content Item", "fr": "Élément de contenu"},
    {
        "label": {"en": "Label", "fr": "Étiquette"},
        "data": {"en": "Data", "fr": "Données"},
        "metadata": {"en": "Metadata", "fr": "Métadonnées"},
    },
)


class ChatContentExtracted(BaseModel):
    id: str = Field(description="Reference to source ChatDocument")
    contents: List[ContentItem] = Field(
        default_factory=list, description="List of content items"
    )


registerModelLabels(
    "ChatContentExtracted",
    {"en": "Extracted Content", "fr": "Contenu extrait"},
    {
        "id": {"en": "Object ID", "fr": "ID de l'objet"},
        "contents": {"en": "Contents", "fr": "Contenus"},
    },
)


class ChatMessage(BaseModel):
    """Messages in chat workflows. User-owned, no mandate context."""
    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"
    )
    documents: List[ChatDocument] = Field(
        default_factory=list, description="Associated documents"
    )
    documentsLabel: Optional[str] = Field(
        None, description="Label for the set of documents"
    )
    message: Optional[str] = Field(None, description="Message content")
    summary: Optional[str] = Field(
        None, description="Short summary of this message for planning/history"
    )
    role: str = Field(description="Role of the message sender")
    status: str = Field(description="Status of the message (first, step, last)")
    sequenceNr: Optional[int] = Field(
        default=0,
        description="Sequence number of the message (set automatically)"
    )
    publishedAt: Optional[float] = Field(
        default=None,
        description="When the message was published (UTC timestamp in seconds)",
    )
    success: Optional[bool] = Field(
        None, description="Whether the message processing was successful"
    )
    actionId: Optional[str] = Field(
        None, description="ID of the action that produced this message"
    )
    actionMethod: Optional[str] = Field(
        None, description="Method of the action that produced this message"
    )
    actionName: Optional[str] = Field(
        None, description="Name of the action that produced this message"
    )
    roundNumber: Optional[int] = Field(None, description="Round number in workflow")
    taskNumber: Optional[int] = Field(None, description="Task number within round")
    actionNumber: Optional[int] = Field(None, description="Action number within task")
    taskProgress: Optional[str] = Field(
        None, description="Task progress status: pending, running, success, fail, retry"
    )
    actionProgress: Optional[str] = Field(
        None, description="Action progress status: pending, running, success, fail"
    )


registerModelLabels(
    "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"},
        "documents": {"en": "Documents", "fr": "Documents"},
        "documentsLabel": {"en": "Documents Label", "fr": "Label des documents"},
        "message": {"en": "Message", "fr": "Message"},
        "summary": {"en": "Summary", "fr": "Résumé"},
        "role": {"en": "Role", "fr": "Rôle"},
        "status": {"en": "Status", "fr": "Statut"},
        "sequenceNr": {"en": "Sequence Number", "fr": "Numéro de séquence"},
        "publishedAt": {"en": "Published At", "fr": "Publié le"},
        "success": {"en": "Success", "fr": "Succès"},
        "actionId": {"en": "Action ID", "fr": "ID de l'action"},
        "actionMethod": {"en": "Action Method", "fr": "Méthode de l'action"},
        "actionName": {"en": "Action Name", "fr": "Nom de l'action"},
        "roundNumber": {"en": "Round Number", "fr": "Numéro de tour"},
        "taskNumber": {"en": "Task Number", "fr": "Numéro de tâche"},
        "actionNumber": {"en": "Action Number", "fr": "Numéro d'action"},
        "taskProgress": {"en": "Task Progress", "fr": "Progression de la tâche"},
        "actionProgress": {"en": "Action Progress", "fr": "Progression de l'action"},
    },
)


class WorkflowModeEnum(str, Enum):
    WORKFLOW_DYNAMIC = "Dynamic"
    WORKFLOW_AUTOMATION = "Automation"
    WORKFLOW_CHATBOT = "Chatbot"
    WORKFLOW_REACT = "React"  # Legacy mode - kept for backward compatibility


registerModelLabels(
    "WorkflowModeEnum",
    {"en": "Workflow Mode", "fr": "Mode de workflow"},
    {
        "WORKFLOW_DYNAMIC": {"en": "Dynamic", "fr": "Dynamique"},
        "WORKFLOW_AUTOMATION": {"en": "Automation", "fr": "Automatisation"},
        "WORKFLOW_CHATBOT": {"en": "Chatbot", "fr": "Chatbot"},
        "WORKFLOW_REACT": {"en": "React (Legacy)", "fr": "React (Hérité)"},
    },
)


class ChatWorkflow(BaseModel):
    """Chat workflow container. User-owned, no mandate context."""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key", json_schema_extra={"frontend_type": "text", "frontend_readonly": True, "frontend_required": False})
    featureInstanceId: Optional[str] = Field(None, description="Feature instance ID for multi-tenancy isolation", json_schema_extra={"frontend_type": "text", "frontend_readonly": True, "frontend_required": False})
    status: str = Field(default="running", description="Current status of the workflow", json_schema_extra={"frontend_type": "select", "frontend_readonly": False, "frontend_required": False, "frontend_options": [
            {"value": "running", "label": {"en": "Running", "fr": "En cours"}},
            {"value": "completed", "label": {"en": "Completed", "fr": "Terminé"}},
            {"value": "stopped", "label": {"en": "Stopped", "fr": "Arrêté"}},
            {"value": "error", "label": {"en": "Error", "fr": "Erreur"}},
        ]})
    name: Optional[str] = Field(None, description="Name of the workflow", json_schema_extra={"frontend_type": "text", "frontend_readonly": False, "frontend_required": True})
    currentRound: int = Field(default=0, description="Current round number", json_schema_extra={"frontend_type": "integer", "frontend_readonly": True, "frontend_required": False})
    currentTask: int = Field(default=0, description="Current task number", json_schema_extra={"frontend_type": "integer", "frontend_readonly": True, "frontend_required": False})
    currentAction: int = Field(default=0, description="Current action number", json_schema_extra={"frontend_type": "integer", "frontend_readonly": True, "frontend_required": False})
    totalTasks: int = Field(default=0, description="Total number of tasks in the workflow", json_schema_extra={"frontend_type": "integer", "frontend_readonly": True, "frontend_required": False})
    totalActions: int = Field(default=0, description="Total number of actions in the workflow", json_schema_extra={"frontend_type": "integer", "frontend_readonly": True, "frontend_required": False})
    lastActivity: float = Field(default_factory=getUtcTimestamp, description="Timestamp of last activity (UTC timestamp in seconds)", json_schema_extra={"frontend_type": "timestamp", "frontend_readonly": True, "frontend_required": False})
    startedAt: float = Field(default_factory=getUtcTimestamp, description="When the workflow started (UTC timestamp in seconds)", json_schema_extra={"frontend_type": "timestamp", "frontend_readonly": True, "frontend_required": False})
    logs: List[ChatLog] = Field(default_factory=list, description="Workflow logs", json_schema_extra={"frontend_type": "text", "frontend_readonly": True, "frontend_required": False})
    messages: List[ChatMessage] = Field(default_factory=list, description="Messages in the workflow", json_schema_extra={"frontend_type": "text", "frontend_readonly": True, "frontend_required": False})
    stats: List[ChatStat] = Field(default_factory=list, description="Workflow statistics list", json_schema_extra={"frontend_type": "text", "frontend_readonly": True, "frontend_required": False})
    tasks: list = Field(default_factory=list, description="List of tasks in the workflow", json_schema_extra={"frontend_type": "text", "frontend_readonly": True, "frontend_required": False})
    workflowMode: WorkflowModeEnum = Field(default=WorkflowModeEnum.WORKFLOW_DYNAMIC, description="Workflow mode selector", json_schema_extra={"frontend_type": "select", "frontend_readonly": False, "frontend_required": False, "frontend_options": [
            {
                "value": WorkflowModeEnum.WORKFLOW_DYNAMIC.value,
                "label": {"en": "Dynamic", "fr": "Dynamique"},
            },
            {
                "value": WorkflowModeEnum.WORKFLOW_AUTOMATION.value,
                "label": {"en": "Automation", "fr": "Automatisation"},
            },
            {
                "value": WorkflowModeEnum.WORKFLOW_CHATBOT.value,
                "label": {"en": "Chatbot", "fr": "Chatbot"},
            },
            {
                "value": WorkflowModeEnum.WORKFLOW_REACT.value,
                "label": {"en": "React (Legacy)", "fr": "React (Hérité)"},
            },
        ]})
    maxSteps: int = Field(default=10, description="Maximum number of iterations in dynamic mode", json_schema_extra={"frontend_type": "integer", "frontend_readonly": False, "frontend_required": False})
    expectedFormats: Optional[List[str]] = Field(None, description="List of expected file format extensions from user request (e.g., ['xlsx', 'pdf']). Extracted during intent analysis.", json_schema_extra={"frontend_type": "text", "frontend_readonly": True, "frontend_required": False})
    
    # Helper methods for execution state management
    def getRoundIndex(self) -> int:
        """Get current round index"""
        return self.currentRound
    
    def getTaskIndex(self) -> int:
        """Get current task index"""
        return self.currentTask
    
    def getActionIndex(self) -> int:
        """Get current action index"""
        return self.currentAction
    
    def incrementRound(self):
        """Increment round when new user input received"""
        self.currentRound += 1
        self.currentTask = 0
        self.currentAction = 0
    
    def incrementTask(self):
        """Increment task when starting new task in current round"""
        self.currentTask += 1
        self.currentAction = 0
    
    def incrementAction(self):
        """Increment action when executing new action in current task"""
        self.currentAction += 1


registerModelLabels(
    "ChatWorkflow",
    {"en": "Chat Workflow", "fr": "Flux de travail de chat"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "featureInstanceId": {"en": "Feature Instance ID", "fr": "ID de l'instance de fonctionnalité"},
        "status": {"en": "Status", "fr": "Statut"},
        "name": {"en": "Name", "fr": "Nom"},
        "currentRound": {"en": "Current Round", "fr": "Tour actuel"},
        "currentTask": {"en": "Current Task", "fr": "Tâche actuelle"},
        "currentAction": {"en": "Current Action", "fr": "Action actuelle"},
        "totalTasks": {"en": "Total Tasks", "fr": "Total des tâches"},
        "totalActions": {"en": "Total Actions", "fr": "Total des actions"},
        "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"},
        "workflowMode": {"en": "Workflow Mode", "fr": "Mode de workflow"},
        "maxSteps": {"en": "Max Steps", "fr": "Étapes max"},
        "expectedFormats": {"en": "Expected Formats", "fr": "Formats attendus"},
    },
)


class UserInputRequest(BaseModel):
    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")
    workflowId: Optional[str] = Field(None, description="Optional ID of the workflow to continue")
    allowedProviders: Optional[List[str]] = Field(None, description="List of allowed AI providers (multiselect)")


registerModelLabels(
    "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"},
        "preferredProvider": {"en": "Preferred Provider", "fr": "Fournisseur préféré"},
    },
)


class ActionDocument(BaseModel):
    """Clear document structure for action results"""

    documentName: str = Field(description="Name of the document")
    documentData: Any = Field(description="Content/data of the document")
    mimeType: str = Field(description="MIME type of the document")
    sourceJson: Optional[Dict[str, Any]] = Field(
        None,
        description="Source JSON structure (preserved when rendering to xlsx/docx/pdf)"
    )
    validationMetadata: Optional[Dict[str, Any]] = Field(
        None,
        description="Action-specific metadata for content validation (e.g., email recipients, attachments, SharePoint paths)"
    )


registerModelLabels(
    "ActionDocument",
    {"en": "Action Document", "fr": "Document d'action"},
    {
        "documentName": {"en": "Document Name", "fr": "Nom du document"},
        "documentData": {"en": "Document Data", "fr": "Données du document"},
        "mimeType": {"en": "MIME Type", "fr": "Type MIME"},
    },
)


class ActionResult(BaseModel):
    """Clean action result with documents as primary output

    IMPORTANT: Action methods should NOT set resultLabel in their return value.
    The resultLabel is managed by the action handler using the action's execResultLabel
    from the action plan. This ensures consistent document routing throughout the workflow.
    """

    success: bool = Field(description="Whether execution succeeded")
    error: Optional[str] = Field(None, description="Error message if failed")
    documents: List[ActionDocument] = Field(
        default_factory=list, description="Document outputs"
    )
    resultLabel: Optional[str] = Field(
        None,
        description="Label for document routing (set by action handler, not by action methods)",
    )

    @classmethod
    def isSuccess(cls, documents: List[ActionDocument] = None) -> "ActionResult":
        return cls(success=True, documents=documents or [])

    @classmethod
    def isFailure(
        cls, error: str, documents: List[ActionDocument] = None
    ) -> "ActionResult":
        return cls(success=False, documents=documents or [], error=error)


registerModelLabels(
    "ActionResult",
    {"en": "Action Result", "fr": "Résultat de l'action"},
    {
        "success": {"en": "Success", "fr": "Succès"},
        "error": {"en": "Error", "fr": "Erreur"},
        "documents": {"en": "Documents", "fr": "Documents"},
        "resultLabel": {"en": "Result Label", "fr": "Étiquette du résultat"},
    },
)


class ActionSelection(BaseModel):
    method: str = Field(description="Method to execute (e.g., web, document, ai)")
    name: str = Field(
        description="Action name within the method (e.g., search, extract)"
    )


registerModelLabels(
    "ActionSelection",
    {"en": "Action Selection", "fr": "Sélection d'action"},
    {
        "method": {"en": "Method", "fr": "Méthode"},
        "name": {"en": "Action Name", "fr": "Nom de l'action"},
    },
)


class ActionParameters(BaseModel):
    parameters: Dict[str, Any] = Field(
        default_factory=dict, description="Parameters to execute the selected action"
    )


registerModelLabels(
    "ActionParameters",
    {"en": "Action Parameters", "fr": "Paramètres d'action"},
    {
        "parameters": {"en": "Parameters", "fr": "Paramètres"},
    },
)


class ObservationPreview(BaseModel):
    name: str = Field(description="Document name or URL label")
    mime: Optional[str] = Field(default=None, description="MIME type or kind (legacy field)")
    snippet: Optional[str] = Field(default=None, description="Short snippet or summary")
    # Extended metadata fields
    mimeType: Optional[str] = Field(default=None, description="MIME type")
    size: Optional[str] = Field(default=None, description="File size")
    created: Optional[str] = Field(default=None, description="Creation timestamp")
    modified: Optional[str] = Field(default=None, description="Modification timestamp")
    typeGroup: Optional[str] = Field(default=None, description="Document type group")
    documentId: Optional[str] = Field(default=None, description="Document ID")
    reference: Optional[str] = Field(default=None, description="Document reference")
    contentSize: Optional[str] = Field(default=None, description="Content size indicator")


registerModelLabels(
    "ObservationPreview",
    {"en": "Observation Preview", "fr": "Aperçu d'observation"},
    {
        "name": {"en": "Name", "fr": "Nom"},
        "mime": {"en": "MIME", "fr": "MIME"},
        "snippet": {"en": "Snippet", "fr": "Extrait"},
    },
)


class Observation(BaseModel):
    success: bool = Field(description="Action execution success flag")
    resultLabel: str = Field(description="Deterministic label for produced documents")
    documentsCount: int = Field(description="Number of produced documents")
    previews: List[ObservationPreview] = Field(
        default_factory=list, description="Compact previews of outputs"
    )
    notes: List[str] = Field(
        default_factory=list, description="Short notes or key facts"
    )
    # Extended fields for enhanced validation
    contentValidation: Optional[Dict[str, Any]] = Field(
        default=None, description="Content validation results"
    )
    contentAnalysis: Optional[Dict[str, Any]] = Field(
        default=None, description="Content analysis results"
    )


registerModelLabels(
    "Observation",
    {"en": "Observation", "fr": "Observation"},
    {
        "success": {"en": "Success", "fr": "Succès"},
        "resultLabel": {"en": "Result Label", "fr": "Étiquette du résultat"},
        "documentsCount": {"en": "Documents Count", "fr": "Nombre de documents"},
        "previews": {"en": "Previews", "fr": "Aperçus"},
        "notes": {"en": "Notes", "fr": "Notes"},
    },
)


class TaskStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


registerModelLabels(
    "TaskStatus",
    {"en": "Task Status", "fr": "Statut de la tâche"},
    {
        "PENDING": {"en": "Pending", "fr": "En attente"},
        "RUNNING": {"en": "Running", "fr": "En cours"},
        "COMPLETED": {"en": "Completed", "fr": "Terminé"},
        "FAILED": {"en": "Failed", "fr": "Échec"},
        "CANCELLED": {"en": "Cancelled", "fr": "Annulé"},
    },
)


class DocumentExchange(BaseModel):
    documentsLabel: str = Field(description="Label for the set of documents")
    documents: List[str] = Field(
        default_factory=list, description="List of document references"
    )


registerModelLabels(
    "DocumentExchange",
    {"en": "Document Exchange", "fr": "Échange de documents"},
    {
        "documentsLabel": {"en": "Documents Label", "fr": "Label des documents"},
        "documents": {"en": "Documents", "fr": "Documents"},
    },
)


class ActionItem(BaseModel):
    id: str = Field(..., description="Action ID")
    execMethod: str = Field(..., description="Method to execute")
    execAction: str = Field(..., description="Action to perform")
    execParameters: Dict[str, Any] = Field(
        default_factory=dict, description="Action parameters"
    )
    execResultLabel: Optional[str] = Field(
        None, description="Label for the set of result documents"
    )
    expectedDocumentFormats: Optional[List[Dict[str, str]]] = Field(
        None, description="Expected document formats (optional)"
    )
    userMessage: Optional[str] = Field(
        None, description="User-friendly message in user's language"
    )
    status: TaskStatus = Field(default=TaskStatus.PENDING, description="Action status")
    error: Optional[str] = Field(None, description="Error message if action failed")
    retryCount: int = Field(default=0, description="Number of retries attempted")
    retryMax: int = Field(default=3, description="Maximum number of retries")
    processingTime: Optional[float] = Field(
        None, description="Processing time in seconds"
    )
    timestamp: float = Field(
        ..., description="When the action was executed (UTC timestamp in seconds)"
    )
    result: Optional[str] = Field(None, description="Result of the action")

    def setSuccess(self, result: str = None) -> None:
        """Set the action as successful with optional result"""
        self.status = TaskStatus.COMPLETED
        self.error = None
        if result is not None:
            self.result = result

    def setError(self, error_message: str) -> None:
        """Set the action as failed with error message"""
        self.status = TaskStatus.FAILED
        self.error = error_message


registerModelLabels(
    "ActionItem",
    {"en": "Task Action", "fr": "Action de tâche"},
    {
        "id": {"en": "Action ID", "fr": "ID de l'action"},
        "execMethod": {"en": "Method", "fr": "Méthode"},
        "execAction": {"en": "Action", "fr": "Action"},
        "execParameters": {"en": "Parameters", "fr": "Paramètres"},
        "execResultLabel": {"en": "Result Label", "fr": "Label du résultat"},
        "expectedDocumentFormats": {
            "en": "Expected Document Formats",
            "fr": "Formats de documents attendus",
        },
        "userMessage": {"en": "User Message", "fr": "Message utilisateur"},
        "status": {"en": "Status", "fr": "Statut"},
        "error": {"en": "Error", "fr": "Erreur"},
        "retryCount": {"en": "Retry Count", "fr": "Nombre de tentatives"},
        "retryMax": {"en": "Max Retries", "fr": "Tentatives max"},
        "processingTime": {"en": "Processing Time", "fr": "Temps de traitement"},
        "timestamp": {"en": "Timestamp", "fr": "Horodatage"},
        "result": {"en": "Result", "fr": "Résultat"},
    },
)


class TaskResult(BaseModel):
    taskId: str = Field(..., description="Task ID")
    status: TaskStatus = Field(default=TaskStatus.PENDING, description="Task status")
    success: bool = Field(..., description="Whether the task was successful")
    feedback: Optional[str] = Field(None, description="Task feedback message")
    error: Optional[str] = Field(None, description="Error message if task failed")


registerModelLabels(
    "TaskResult",
    {"en": "Task Result", "fr": "Résultat de tâche"},
    {
        "taskId": {"en": "Task ID", "fr": "ID de la tâche"},
        "status": {"en": "Status", "fr": "Statut"},
        "success": {"en": "Success", "fr": "Succès"},
        "feedback": {"en": "Feedback", "fr": "Retour"},
        "error": {"en": "Error", "fr": "Erreur"},
    },
)


class TaskItem(BaseModel):
    id: str = Field(..., description="Task ID")
    workflowId: str = Field(..., description="Workflow ID")
    userInput: str = Field(..., description="User input that triggered the task")
    status: TaskStatus = Field(default=TaskStatus.PENDING, description="Task status")
    error: Optional[str] = Field(None, description="Error message if task failed")
    startedAt: Optional[float] = Field(
        None, description="When the task started (UTC timestamp in seconds)"
    )
    finishedAt: Optional[float] = Field(
        None, description="When the task finished (UTC timestamp in seconds)"
    )
    actionList: List[ActionItem] = Field(
        default_factory=list, description="List of actions to execute"
    )
    retryCount: int = Field(default=0, description="Number of retries attempted")
    retryMax: int = Field(default=3, description="Maximum number of retries")
    rollbackOnFailure: bool = Field(
        default=True, description="Whether to rollback on failure"
    )
    dependencies: List[str] = Field(
        default_factory=list, description="List of task IDs this task depends on"
    )
    feedback: Optional[str] = Field(None, description="Task feedback message")
    processingTime: Optional[float] = Field(
        None, description="Total processing time in seconds"
    )
    resultLabels: Optional[Dict[str, Any]] = Field(
        default_factory=dict, description="Map of result labels to their values"
    )


registerModelLabels(
    "TaskItem",
    {"en": "Task", "fr": "Tâche"},
    {
        "id": {"en": "Task ID", "fr": "ID de la tâche"},
        "workflowId": {"en": "Workflow ID", "fr": "ID du workflow"},
        "userInput": {"en": "User Input", "fr": "Entrée utilisateur"},
        "status": {"en": "Status", "fr": "Statut"},
        "error": {"en": "Error", "fr": "Erreur"},
        "startedAt": {"en": "Started At", "fr": "Démarré à"},
        "finishedAt": {"en": "Finished At", "fr": "Terminé à"},
        "actionList": {"en": "Actions", "fr": "Actions"},
        "retryCount": {"en": "Retry Count", "fr": "Nombre de tentatives"},
        "retryMax": {"en": "Max Retries", "fr": "Tentatives max"},
        "processingTime": {"en": "Processing Time", "fr": "Temps de traitement"},
    },
)


class TaskStep(BaseModel):
    id: str
    objective: str
    dependencies: Optional[list[str]] = Field(default_factory=list)
    successCriteria: Optional[list[str]] = Field(default_factory=list)
    estimatedComplexity: Optional[str] = None
    userMessage: Optional[str] = Field(
        None, description="User-friendly message in user's language"
    )
    # Format details extracted from intent analysis
    dataType: Optional[str] = Field(
        None, description="Expected data type (text, numbers, documents, etc.)"
    )
    expectedFormats: Optional[List[str]] = Field(
        None, description="Expected output file format extensions (e.g., ['docx', 'pdf', 'xlsx']). Use actual file extensions, not conceptual terms."
    )
    qualityRequirements: Optional[Dict[str, Any]] = Field(
        None, description="Quality requirements and constraints"
    )


registerModelLabels(
    "TaskStep",
    {"en": "Task Step", "fr": "Étape de tâche"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "objective": {"en": "Objective", "fr": "Objectif"},
        "dependencies": {"en": "Dependencies", "fr": "Dépendances"},
        "successCriteria": {"en": "Success Criteria", "fr": "Critères de succès"},
        "estimatedComplexity": {
            "en": "Estimated Complexity",
            "fr": "Complexité estimée",
        },
        "userMessage": {"en": "User Message", "fr": "Message utilisateur"},
        "expectedFormats": {"en": "Expected Formats", "fr": "Formats attendus"},
    },
)


class TaskHandover(BaseModel):
    taskId: str = Field(description="Target task ID")
    sourceTask: Optional[str] = Field(None, description="Source task ID")
    inputDocuments: List[DocumentExchange] = Field(
        default_factory=list, description="Available input documents"
    )
    outputDocuments: List[DocumentExchange] = Field(
        default_factory=list, description="Produced output documents"
    )
    context: Dict[str, Any] = Field(default_factory=dict, description="Task context")
    previousResults: List[str] = Field(
        default_factory=list, description="Previous result summaries"
    )
    improvements: List[str] = Field(
        default_factory=list, description="Improvement suggestions"
    )
    workflowSummary: Optional[str] = Field(
        None, description="Summarized workflow context"
    )
    messageHistory: List[str] = Field(
        default_factory=list, description="Key message summaries"
    )
    timestamp: float = Field(
        ..., description="When the handover was created (UTC timestamp in seconds)"
    )
    handoverType: str = Field(
        default="task", description="Type of handover: task, phase, or workflow"
    )


registerModelLabels(
    "TaskHandover",
    {"en": "Task Handover", "fr": "Transfert de tâche"},
    {
        "taskId": {"en": "Task ID", "fr": "ID de la tâche"},
        "sourceTask": {"en": "Source Task", "fr": "Tâche source"},
        "inputDocuments": {"en": "Input Documents", "fr": "Documents d'entrée"},
        "outputDocuments": {"en": "Output Documents", "fr": "Documents de sortie"},
        "context": {"en": "Context", "fr": "Contexte"},
        "previousResults": {"en": "Previous Results", "fr": "Résultats précédents"},
        "improvements": {"en": "Improvements", "fr": "Améliorations"},
        "workflowSummary": {"en": "Workflow Summary", "fr": "Résumé du workflow"},
        "messageHistory": {"en": "Message History", "fr": "Historique des messages"},
        "timestamp": {"en": "Timestamp", "fr": "Horodatage"},
        "handoverType": {"en": "Handover Type", "fr": "Type de transfert"},
    },
)


class TaskContext(BaseModel):
    taskStep: TaskStep
    workflow: Optional[ChatWorkflow] = None
    workflowId: Optional[str] = None
    availableDocuments: Optional[str] = "No documents available"
    availableConnections: Optional[list[str]] = Field(default_factory=list)
    previousResults: Optional[list[str]] = Field(default_factory=list)
    previousHandover: Optional[TaskHandover] = None
    improvements: Optional[list[str]] = Field(default_factory=list)
    retryCount: Optional[int] = 0
    previousActionResults: Optional[list] = Field(default_factory=list)
    previousReviewResult: Optional[dict] = None
    isRegeneration: Optional[bool] = False
    failurePatterns: Optional[list[str]] = Field(default_factory=list)
    failedActions: Optional[list] = Field(default_factory=list)
    successfulActions: Optional[list] = Field(default_factory=list)
    executedActions: Optional[list] = Field(default_factory=list, description="List of executed actions with action name, parameters, and step number")
    criteriaProgress: Optional[dict] = None
    
    # Stage 2 context fields (NEW)
    actionObjective: Optional[str] = Field(None, description="Objective for current action")
    parametersContext: Optional[str] = Field(None, description="Context for parameter generation")
    learnings: Optional[list[str]] = Field(default_factory=list, description="Learnings from previous actions")
    stage1Selection: Optional[dict] = Field(None, description="Stage 1 selection data")
    nextActionGuidance: Optional[Dict[str, Any]] = Field(None, description="Guidance for the next action from previous refinement")
    
    def updateFromSelection(self, selection: Any):
        """Update context from Stage 1 selection
        
        Args:
            selection: ActionDefinition instance from Stage 1
        """
        from modules.datamodels.datamodelWorkflow import ActionDefinition
        
        if isinstance(selection, ActionDefinition):
            self.actionObjective = selection.actionObjective
            self.parametersContext = selection.parametersContext
            self.learnings = selection.learnings if selection.learnings else []
            self.stage1Selection = selection.model_dump()

    def getDocumentReferences(self) -> List[str]:
        docs = []
        if self.previousHandover:
            for doc_exchange in self.previousHandover.inputDocuments:
                docs.extend(doc_exchange.documents)
        return list(set(docs))

    def addImprovement(self, improvement: str) -> None:
        if improvement not in (self.improvements or []):
            if self.improvements is None:
                self.improvements = []
            self.improvements.append(improvement)


class ReviewContext(BaseModel):
    taskStep: TaskStep
    taskActions: Optional[list] = Field(default_factory=list)
    actionResults: Optional[list] = Field(default_factory=list)
    stepResult: Optional[dict] = Field(default_factory=dict)
    workflowId: Optional[str] = None
    previousResults: Optional[list[str]] = Field(default_factory=list)


class ReviewResult(BaseModel):
    status: str
    reason: Optional[str] = None
    improvements: Optional[list[str]] = Field(default_factory=list)
    qualityScore: Optional[float] = Field(default=5.0, description="Quality score (0-10)")
    missingOutputs: Optional[list[str]] = Field(default_factory=list)
    metCriteria: Optional[list[str]] = Field(default_factory=list)
    unmetCriteria: Optional[list[str]] = Field(default_factory=list)
    confidence: Optional[float] = 0.5
    userMessage: Optional[str] = Field(
        None, description="User-friendly message in user's language"
    )
    # NEW: Concrete next action guidance (when status is "continue")
    nextAction: Optional[str] = Field(
        None, description="Specific action to execute next (e.g., 'ai.convert', 'ai.process', 'ai.reformat')"
    )
    nextActionParameters: Optional[Dict[str, Any]] = Field(
        None, description="Parameters for the next action (e.g., {'fromFormat': 'json', 'toFormat': 'csv'})"
    )
    nextActionObjective: Optional[str] = Field(
        None, description="What this specific action will achieve"
    )


registerModelLabels(
    "ReviewResult",
    {"en": "Review Result", "fr": "Résultat de l'évaluation"},
    {
        "status": {"en": "Status", "fr": "Statut"},
        "reason": {"en": "Reason", "fr": "Raison"},
        "improvements": {"en": "Improvements", "fr": "Améliorations"},
        "qualityScore": {"en": "Quality Score", "fr": "Score de qualité"},
        "missingOutputs": {"en": "Missing Outputs", "fr": "Sorties manquantes"},
        "metCriteria": {"en": "Met Criteria", "fr": "Critères respectés"},
        "unmetCriteria": {"en": "Unmet Criteria", "fr": "Critères non respectés"},
        "confidence": {"en": "Confidence", "fr": "Confiance"},
        "userMessage": {"en": "User Message", "fr": "Message utilisateur"},
    },
)


class TaskPlan(BaseModel):
    overview: str
    tasks: list[TaskStep]
    userMessage: Optional[str] = Field(
        None, description="Overall user-friendly message for the task plan"
    )


registerModelLabels(
    "TaskPlan",
    {"en": "Task Plan", "fr": "Plan de tâches"},
    {
        "overview": {"en": "Overview", "fr": "Aperçu"},
        "tasks": {"en": "Tasks", "fr": "Tâches"},
        "userMessage": {"en": "User Message", "fr": "Message utilisateur"},
    },
)

# Forward references resolved automatically since ChatWorkflow is defined above


class PromptPlaceholder(BaseModel):
    label: str
    content: str
    summaryAllowed: bool = Field(
        default=False,
        description="Whether host may summarize content before sending to AI",
    )


registerModelLabels(
    "PromptPlaceholder",
    {"en": "Prompt Placeholder", "fr": "Espace réservé d'invite"},
    {
        "label": {"en": "Label", "fr": "Libellé"},
        "content": {"en": "Content", "fr": "Contenu"},
        "summaryAllowed": {"en": "Summary Allowed", "fr": "Résumé autorisé"},
    },
)


class PromptBundle(BaseModel):
    prompt: str
    placeholders: List[PromptPlaceholder] = Field(default_factory=list)


registerModelLabels(
    "PromptBundle",
    {"en": "Prompt Bundle", "fr": "Lot d'invite"},
    {
        "prompt": {"en": "Prompt", "fr": "Invite"},
        "placeholders": {"en": "Placeholders", "fr": "Espaces réservés"},
    },
)
