"""
Workflow execution models for action definitions, AI responses, and workflow-level structures.
"""

from typing import Dict, Any, List, Optional, TYPE_CHECKING
from pydantic import BaseModel, Field
from modules.shared.attributeUtils import registerModelLabels
from modules.shared.jsonUtils import extractJsonString, tryParseJson, repairBrokenJson

# Import DocumentReferenceList at runtime (needed for ActionDefinition)
from modules.datamodels.datamodelDocref import DocumentReferenceList

# Forward references for circular imports (use string annotations)
if TYPE_CHECKING:
    from modules.datamodels.datamodelChat import ChatDocument, ActionResult
    from modules.datamodels.datamodelExtraction import ExtractionOptions


class ActionDefinition(BaseModel):
    """Action definition with selection and parameters from planning phase"""
    
    # Core action selection (Stage 1)
    action: str = Field(description="Compound action name (method.action)")
    actionObjective: str = Field(description="Objective for this action")
    userMessage: Optional[str] = Field(
        None,
        description="User-friendly message in user's language explaining what this action will do (generated by AI in prompts)"
    )
    parametersContext: Optional[str] = Field(
        None,
        description="Context for parameter generation"
    )
    learnings: List[str] = Field(
        default_factory=list,
        description="Learnings from previous actions"
    )
    
    # Resources (ALWAYS defined in Stage 1 if action needs them)
    documentList: Optional[DocumentReferenceList] = Field(
        None,
        description="Document references (ALWAYS defined in Stage 1 if action needs documents)"
    )
    connectionReference: Optional[str] = Field(
        None,
        description="Connection reference (ALWAYS defined in Stage 1 if action needs connection)"
    )
    
    # Parameters (may be defined in Stage 1 OR Stage 2, depending on action and actionObjective)
    parameters: Optional[Dict[str, Any]] = Field(
        None,
        description="Action-specific parameters (generated in Stage 2 for complex actions, or inferred from actionObjective for simple actions)"
    )
    
    def hasParameters(self) -> bool:
        """Check if parameters have been generated (Stage 2 complete or inferred)"""
        return self.parameters is not None
    
    def needsStage2(self) -> bool:
        """Determine if Stage 2 parameter generation is needed (generic, deterministic check)
        
        Generic logic (works for any action, dynamically added or removed):
        - If parameters are already set → Stage 2 not needed
        - If parameters are None → Stage 2 needed (to generate parameters from actionObjective and context)
        
        Note: Stage 1 always defines documentList and connectionReference if the action needs them.
        Stage 2 only generates the action-specific parameters dictionary.
        """
        # Generic check: if parameters are not set, Stage 2 is needed
        return self.parameters is None
    
    def updateFromStage1StringReferences(self, stringRefs: Optional[List[str]], connectionRef: Optional[str]):
        """Update documentList and connectionReference from Stage 1 string references
        
        Called when Stage 1 AI returns string references that need to be converted to typed models.
        """
        if stringRefs:
            self.documentList = DocumentReferenceList.from_string_list(stringRefs)
        if connectionRef:
            self.connectionReference = connectionRef


class AiResponseMetadata(BaseModel):
    """Metadata for AI response (varies by operation type)."""
    
    # Document Generation Metadata
    title: Optional[str] = Field(None, description="Document title")
    filename: Optional[str] = Field(None, description="Document filename")
    
    # Operation-Specific Metadata
    operationType: Optional[str] = Field(None, description="Type of operation performed")
    schemaVersion: Optional[str] = Field(None, description="Schema version (e.g., 'parameters_v1')", alias="schema")
    extractionMethod: Optional[str] = Field(None, description="Method used for extraction")
    sourceDocuments: Optional[List[str]] = Field(None, description="Source document references")
    
    # Additional metadata (for extensibility)
    additionalData: Optional[Dict[str, Any]] = Field(None, description="Additional operation-specific metadata")


class DocumentData(BaseModel):
    """Single document in response"""
    documentName: str = Field(description="Document name")
    documentData: Any = Field(description="Document data (can be str, bytes, dict, etc.)")
    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)"
    )


class ExtractContentParameters(BaseModel):
    """Parameters for extraction action.
    
    This model is defined together with the `methodAi.extractContent()` action function.
    All action parameter models follow this pattern: defined in the same module as the action.
    However, since this is a workflow-level model used across the system, it's defined here.
    """
    documentList: DocumentReferenceList = Field(description="Document references to extract content from")
    extractionOptions: Optional[Any] = Field(  # ExtractionOptions - forward reference
        None,
        description="Extraction options (determined dynamically based on task and document characteristics)"
    )


class AiResponse(BaseModel):
    """Unified response from all AI calls (planning, text, documents)"""
    
    content: str = Field(description="Response content (JSON string for planning, text for analysis, unified JSON for documents)")
    metadata: Optional[AiResponseMetadata] = Field(
        None,
        description="Response metadata (varies by operation type)"
    )
    documents: Optional[List[DocumentData]] = Field(
        None,
        description="Generated documents (only for document generation operations)"
    )
    
    def toJson(self) -> Dict[str, Any]:
        """
        Convert AI response content to JSON using enhanced stabilizing failsafe conversion methods.
        Centralizes AI result to JSON conversion in one place.
        
        Uses methods from jsonUtils:
        - tryParseJson() - Safe parsing with error handling
        - repairBrokenJson() - Repairs broken/incomplete JSON
        - extractJsonString() - Extracts JSON from text with code fences
        
        Returns:
            Dict containing the parsed JSON content, or a safe fallback structure if parsing fails.
            - If content is valid JSON dict: returns the dict directly
            - If content is valid JSON list: wraps in {"data": [...]}
            - If content is broken JSON: attempts repair using repairBrokenJson()
            - If all parsing fails: returns {"content": "...", "parseError": True}
        """
        # If content is already a dict, return it directly
        if isinstance(self.content, dict):
            return self.content
        
        # If content is already a list, wrap it
        if isinstance(self.content, list):
            return {"data": self.content}
        
        # Convert to string if needed
        contentStr = str(self.content) if not isinstance(self.content, str) else self.content
        
        # First, try to extract JSON from text (handles code fences, etc.)
        extractedJson = extractJsonString(contentStr)
        
        # Try to parse as JSON (returns tuple: obj, error, cleaned_str)
        parsedJson, parseError, _ = tryParseJson(extractedJson)
        
        if parsedJson is not None and parseError is None:
            # If it's a dict, return directly
            if isinstance(parsedJson, dict):
                return parsedJson
            # If it's a list, wrap in dict
            elif isinstance(parsedJson, list):
                return {"data": parsedJson}
        
        # Try to repair broken JSON
        repairedJson = repairBrokenJson(contentStr)
        if repairedJson:
            # repairBrokenJson returns Optional[Dict[str, Any]] - always a dict or None
            if isinstance(repairedJson, dict):
                return repairedJson
        
        # All parsing failed - return safe fallback
        contentStr = str(self.content) if not isinstance(self.content, str) else self.content
        return {"content": contentStr, "parseError": True}


# Workflow-level models

class RequestContext(BaseModel):
    """Normalized request context from user input"""
    
    originalPrompt: str = Field(description="Original user prompt")
    documents: List[Any] = Field(  # ChatDocument - forward reference
        default_factory=list,
        description="Documents provided by user"
    )
    userLanguage: str = Field(description="User's language")
    detectedComplexity: str = Field(
        description="Complexity level: simple, moderate, complex"
    )
    requiresDocuments: bool = Field(default=False, description="Whether request requires documents")
    requiresWebResearch: bool = Field(default=False, description="Whether request requires web research")
    requiresAnalysis: bool = Field(default=False, description="Whether request requires analysis")
    expectedOutputFormat: Optional[str] = Field(None, description="Expected output format")
    expectedOutputType: Optional[str] = Field(None, description="Expected output type: answer, document, analysis")


class UnderstandingResult(BaseModel):
    """Result from initial understanding phase (combined AI call)"""
    
    parameters: Dict[str, Any] = Field(
        default_factory=dict,
        description="Basic parameters (language, format, detail level)"
    )
    intention: Dict[str, Any] = Field(
        default_factory=dict,
        description="User intention (primaryGoal, secondaryGoals, intentionType)"
    )
    context: Dict[str, Any] = Field(
        default_factory=dict,
        description="Extracted context (topics, requirements, constraints)"
    )
    documentReferences: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Document references with purpose and relevance"
    )
    tasks: List["TaskDefinition"] = Field(  # Forward reference
        default_factory=list,
        description="Task definitions with deliverables"
    )


class TaskDefinition(BaseModel):
    """Task definition from understanding phase"""
    
    id: str = Field(description="Task identifier")
    objective: str = Field(description="Task objective")
    deliverable: Dict[str, Any] = Field(
        description="Deliverable specification (type, format, style, detailLevel)"
    )
    requiresWebResearch: bool = Field(default=False, description="Whether task requires web research")
    requiresDocumentAnalysis: bool = Field(default=False, description="Whether task requires document analysis")
    requiresContentGeneration: bool = Field(default=True, description="Whether task requires content generation")
    requiredDocuments: List[str] = Field(
        default_factory=list,
        description="Document references needed for this task"
    )
    extractionOptions: Optional[Any] = Field(  # ExtractionOptions - forward reference
        None,
        description="Extraction options for document processing (determined dynamically based on task and document characteristics)"
    )


class TaskResult(BaseModel):
    """Result from task execution"""
    
    taskId: str = Field(description="Task identifier")
    actionResult: Any = Field(description="ActionResult from task execution")  # ActionResult - forward reference


# Register model labels for UI
registerModelLabels(
    "RequestContext",
    {"en": "Request Context", "fr": "Contexte de la demande"},
    {
        "originalPrompt": {"en": "Original Prompt", "fr": "Invite originale"},
        "documents": {"en": "Documents", "fr": "Documents"},
        "userLanguage": {"en": "User Language", "fr": "Langue de l'utilisateur"},
        "detectedComplexity": {"en": "Detected Complexity", "fr": "Complexité détectée"},
        "requiresDocuments": {"en": "Requires Documents", "fr": "Nécessite des documents"},
        "requiresWebResearch": {"en": "Requires Web Research", "fr": "Nécessite une recherche web"},
        "requiresAnalysis": {"en": "Requires Analysis", "fr": "Nécessite une analyse"},
        "expectedOutputFormat": {"en": "Expected Output Format", "fr": "Format de sortie attendu"},
        "expectedOutputType": {"en": "Expected Output Type", "fr": "Type de sortie attendu"},
    },
)

registerModelLabels(
    "UnderstandingResult",
    {"en": "Understanding Result", "fr": "Résultat de compréhension"},
    {
        "parameters": {"en": "Parameters", "fr": "Paramètres"},
        "intention": {"en": "Intention", "fr": "Intention"},
        "context": {"en": "Context", "fr": "Contexte"},
        "documentReferences": {"en": "Document References", "fr": "Références de documents"},
        "tasks": {"en": "Tasks", "fr": "Tâches"},
    },
)

registerModelLabels(
    "TaskDefinition",
    {"en": "Task Definition", "fr": "Définition de tâche"},
    {
        "id": {"en": "Task ID", "fr": "ID de la tâche"},
        "objective": {"en": "Objective", "fr": "Objectif"},
        "deliverable": {"en": "Deliverable", "fr": "Livrable"},
        "requiresWebResearch": {"en": "Requires Web Research", "fr": "Nécessite une recherche web"},
        "requiresDocumentAnalysis": {"en": "Requires Document Analysis", "fr": "Nécessite une analyse de documents"},
        "requiresContentGeneration": {"en": "Requires Content Generation", "fr": "Nécessite une génération de contenu"},
        "requiredDocuments": {"en": "Required Documents", "fr": "Documents requis"},
        "extractionOptions": {"en": "Extraction Options", "fr": "Options d'extraction"},
    },
)

registerModelLabels(
    "TaskResult",
    {"en": "Task Result", "fr": "Résultat de tâche"},
    {
        "taskId": {"en": "Task ID", "fr": "ID de la tâche"},
        "actionResult": {"en": "Action Result", "fr": "Résultat de l'action"},
    },
)

registerModelLabels(
    "RequestContext",
    {"en": "Request Context", "fr": "Contexte de la demande"},
    {
        "originalPrompt": {"en": "Original Prompt", "fr": "Invite originale"},
        "documents": {"en": "Documents", "fr": "Documents"},
        "userLanguage": {"en": "User Language", "fr": "Langue de l'utilisateur"},
        "detectedComplexity": {"en": "Detected Complexity", "fr": "Complexité détectée"},
        "requiresDocuments": {"en": "Requires Documents", "fr": "Nécessite des documents"},
        "requiresWebResearch": {"en": "Requires Web Research", "fr": "Nécessite une recherche web"},
        "requiresAnalysis": {"en": "Requires Analysis", "fr": "Nécessite une analyse"},
        "expectedOutputFormat": {"en": "Expected Output Format", "fr": "Format de sortie attendu"},
        "expectedOutputType": {"en": "Expected Output Type", "fr": "Type de sortie attendu"},
    },
)

registerModelLabels(
    "UnderstandingResult",
    {"en": "Understanding Result", "fr": "Résultat de compréhension"},
    {
        "parameters": {"en": "Parameters", "fr": "Paramètres"},
        "intention": {"en": "Intention", "fr": "Intention"},
        "context": {"en": "Context", "fr": "Contexte"},
        "documentReferences": {"en": "Document References", "fr": "Références de documents"},
        "tasks": {"en": "Tasks", "fr": "Tâches"},
    },
)

registerModelLabels(
    "TaskDefinition",
    {"en": "Task Definition", "fr": "Définition de tâche"},
    {
        "id": {"en": "Task ID", "fr": "ID de la tâche"},
        "objective": {"en": "Objective", "fr": "Objectif"},
        "deliverable": {"en": "Deliverable", "fr": "Livrable"},
        "requiresWebResearch": {"en": "Requires Web Research", "fr": "Nécessite une recherche web"},
        "requiresDocumentAnalysis": {"en": "Requires Document Analysis", "fr": "Nécessite une analyse de documents"},
        "requiresContentGeneration": {"en": "Requires Content Generation", "fr": "Nécessite une génération de contenu"},
        "requiredDocuments": {"en": "Required Documents", "fr": "Documents requis"},
        "extractionOptions": {"en": "Extraction Options", "fr": "Options d'extraction"},
    },
)

registerModelLabels(
    "TaskResult",
    {"en": "Task Result", "fr": "Résultat de tâche"},
    {
        "taskId": {"en": "Task ID", "fr": "ID de la tâche"},
        "actionResult": {"en": "Action Result", "fr": "Résultat de l'action"},
    },
)

# Register model labels for UI
registerModelLabels(
    "ActionDefinition",
    {"en": "Action Definition", "fr": "Définition d'action"},
    {
        "action": {"en": "Action", "fr": "Action"},
        "actionObjective": {"en": "Action Objective", "fr": "Objectif de l'action"},
        "parametersContext": {"en": "Parameters Context", "fr": "Contexte des paramètres"},
        "learnings": {"en": "Learnings", "fr": "Apprentissages"},
        "documentList": {"en": "Document List", "fr": "Liste de documents"},
        "connectionReference": {"en": "Connection Reference", "fr": "Référence de connexion"},
        "parameters": {"en": "Parameters", "fr": "Paramètres"},
    },
)

registerModelLabels(
    "AiResponse",
    {"en": "AI Response", "fr": "Réponse IA"},
    {
        "content": {"en": "Content", "fr": "Contenu"},
        "metadata": {"en": "Metadata", "fr": "Métadonnées"},
        "documents": {"en": "Documents", "fr": "Documents"},
    },
)

registerModelLabels(
    "AiResponseMetadata",
    {"en": "AI Response Metadata", "fr": "Métadonnées de réponse IA"},
    {
        "title": {"en": "Title", "fr": "Titre"},
        "filename": {"en": "Filename", "fr": "Nom de fichier"},
        "operationType": {"en": "Operation Type", "fr": "Type d'opération"},
        "schemaVersion": {"en": "Schema Version", "fr": "Version du schéma"},
        "extractionMethod": {"en": "Extraction Method", "fr": "Méthode d'extraction"},
        "sourceDocuments": {"en": "Source Documents", "fr": "Documents sources"},
    },
)

registerModelLabels(
    "DocumentData",
    {"en": "Document Data", "fr": "Données de document"},
    {
        "documentName": {"en": "Document Name", "fr": "Nom du document"},
        "documentData": {"en": "Document Data", "fr": "Données du document"},
        "mimeType": {"en": "MIME Type", "fr": "Type MIME"},
    },
)

registerModelLabels(
    "RequestContext",
    {"en": "Request Context", "fr": "Contexte de requête"},
    {
        "originalPrompt": {"en": "Original Prompt", "fr": "Invite originale"},
        "documents": {"en": "Documents", "fr": "Documents"},
        "userLanguage": {"en": "User Language", "fr": "Langue de l'utilisateur"},
        "detectedComplexity": {"en": "Detected Complexity", "fr": "Complexité détectée"},
        "requiresDocuments": {"en": "Requires Documents", "fr": "Nécessite des documents"},
        "requiresWebResearch": {"en": "Requires Web Research", "fr": "Nécessite une recherche web"},
        "requiresAnalysis": {"en": "Requires Analysis", "fr": "Nécessite une analyse"},
    },
)

registerModelLabels(
    "UnderstandingResult",
    {"en": "Understanding Result", "fr": "Résultat de compréhension"},
    {
        "parameters": {"en": "Parameters", "fr": "Paramètres"},
        "intention": {"en": "Intention", "fr": "Intention"},
        "context": {"en": "Context", "fr": "Contexte"},
        "documentReferences": {"en": "Document References", "fr": "Références de documents"},
        "tasks": {"en": "Tasks", "fr": "Tâches"},
    },
)

registerModelLabels(
    "TaskDefinition",
    {"en": "Task Definition", "fr": "Définition de tâche"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "objective": {"en": "Objective", "fr": "Objectif"},
        "deliverable": {"en": "Deliverable", "fr": "Livrable"},
        "requiresWebResearch": {"en": "Requires Web Research", "fr": "Nécessite une recherche web"},
        "requiresDocumentAnalysis": {"en": "Requires Document Analysis", "fr": "Nécessite une analyse de document"},
        "requiresContentGeneration": {"en": "Requires Content Generation", "fr": "Nécessite une génération de contenu"},
        "requiredDocuments": {"en": "Required Documents", "fr": "Documents requis"},
        "extractionOptions": {"en": "Extraction Options", "fr": "Options d'extraction"},
    },
)

registerModelLabels(
    "TaskResult",
    {"en": "Task Result", "fr": "Résultat de tâche"},
    {
        "taskId": {"en": "Task ID", "fr": "ID de tâche"},
        "actionResult": {"en": "Action Result", "fr": "Résultat d'action"},
    },
)