"""
Refactored WorkflowManager class for the Agentservice (continued).
"""

import os
import logging
import asyncio
import uuid
from datetime import datetime
from typing import List, Dict, Any, Optional, Tuple, Union

logger = logging.getLogger(__name__)

class WorkflowManager:
    # Previous code is in the first part
    
    async def list_workflows(self, mandate_id: int = None, user_id: int = None) -> List[Dict[str, Any]]:
        """
        List all available workflows.
        
        Args:
            mandate_id: Optional mandate ID for filtering
            user_id: Optional user ID for filtering
            
        Returns:
            List of workflow summaries
        """
        workflows = []
        
        # Load from database if available
        if self.lucydom_interface:
            try:
                # Get all workflows for the user
                if user_id is not None:
                    user_workflows = self.lucydom_interface.get_workflows_by_user(user_id)
                else:
                    user_workflows = self.lucydom_interface.get_all_workflows()
                
                # Filter by mandate if specified
                if mandate_id is not None:
                    user_workflows = [wf for wf in user_workflows if wf.get("mandate_id") == mandate_id]
                
                # Create workflow summaries
                for workflow in user_workflows:
                    summary = {
                        "id": workflow.get("id"),
                        "name": workflow.get("name", f"Workflow {workflow.get('id')}"),
                        "status": workflow.get("status"),
                        "started_at": workflow.get("started_at"),
                        "last_activity": workflow.get("last_activity"),
                        "completed_at": workflow.get("completed_at")
                    }
                    
                    # Add message count if available
                    messages = self.lucydom_interface.get_workflow_messages(workflow.get("id"))
                    if messages:
                        summary["message_count"] = len(messages)
                    
                    workflows.append(summary)
                
                logger.info(f"Loaded {len(workflows)} workflows from database")
                
                # Sort by last activity (newest first)
                return sorted(workflows, key=lambda w: w.get("last_activity", ""), reverse=True)
            
            except Exception as e:
                logger.error(f"Error retrieving workflows from database: {str(e)}")
        
        # Load from files if no database or error occurred
        try:
            for filename in os.listdir(self.results_dir):
                if filename.startswith("workflow_") and filename.endswith(".json"):
                    workflow_path = os.path.join(self.results_dir, filename)
                    
                    try:
                        import json
                        with open(workflow_path, 'r', encoding='utf-8') as f:
                            workflow = json.load(f)
                        
                        # Check if mandate and user ID match filters
                        if mandate_id is not None and workflow.get("mandate_id") != mandate_id:
                            continue
                        
                        if user_id is not None and workflow.get("user_id") != user_id:
                            continue
                        
                        # Create workflow summary
                        summary = {
                            "id": workflow.get("id"),
                            "name": workflow.get("name", f"Workflow {workflow.get('id')}"),
                            "status": workflow.get("status"),
                            "started_at": workflow.get("started_at"),
                            "last_activity": workflow.get("last_activity"),
                            "message_count": len(workflow.get("messages", []))
                        }
                        
                        workflows.append(summary)
                    except Exception as e:
                        logger.error(f"Error loading workflow file {filename}: {str(e)}")
            
            logger.info(f"Loaded {len(workflows)} workflows from files")
            
            # Sort by last activity (newest first)
            return sorted(workflows, key=lambda w: w.get("last_activity", ""), reverse=True)
        
        except Exception as e:
            logger.error(f"Error listing workflows: {str(e)}")
            return []
    
    async def delete_workflow(self, workflow_id: str) -> bool:
        """
        Delete a workflow.
        
        Args:
            workflow_id: ID of the workflow
            
        Returns:
            True on success, False if workflow not found
        """
        # Remove from memory
        if workflow_id in self.workflows:
            del self.workflows[workflow_id]
        
        # Delete from database
        if self.lucydom_interface:
            try:
                db_success = self.lucydom_interface.delete_workflow(workflow_id)
                logger.info(f"Workflow {workflow_id} deleted from database: {db_success}")
            except Exception as e:
                logger.error(f"Error deleting workflow {workflow_id} from database: {str(e)}")
        
        # Delete file
        workflow_path = os.path.join(self.results_dir, f"workflow_{workflow_id}.json")
        
        try:
            if os.path.exists(workflow_path):
                os.remove(workflow_path)
                logger.info(f"Workflow {workflow_id} deleted from file: {workflow_path}")
                return True
            else:
                logger.warning(f"Workflow {workflow_id} not found: {workflow_path}")
                return False
        except Exception as e:
            logger.error(f"Error deleting workflow file {workflow_id}: {str(e)}")
            return False
    
    def _initialize_workflow(self, workflow_id: str) -> Dict[str, Any]:
        """
        Initialize a new workflow.
        
        Args:
            workflow_id: ID of the workflow
            
        Returns:
            The initialized workflow object
        """
        current_time = datetime.now().isoformat()
        
        # Create complete workflow object according to the data model
        workflow = {
            "id": workflow_id,
            "name": f"Workflow {workflow_id}",
            "mandate_id": self.mandate_id,
            "user_id": self.user_id,
            "status": "running",
            "started_at": current_time,
            "last_activity": current_time,
            "current_round": 1,
            
            # Complete statistics structure according to DataStats model
            "data_stats": {
                "total_processing_time": 0.0,
                "total_token_count": 0,
                "total_bytes_sent": 0,
                "total_bytes_received": 0
            },
            
            # Empty arrays for messages and logs
            "messages": [],
            "logs": []
        }
        
        # Log entry for workflow start
        self._add_log(workflow, "Workflow started", "info", "workflow", "Workflow Management")
        
        # Save workflow to database
        if self.lucydom_interface:
            try:
                # Direct save of the complete workflow object
                self.lucydom_interface.save_workflow_state(workflow)
                logger.info(f"Workflow {workflow_id} created in database")
            except Exception as e:
                logger.error(f"Error creating workflow {workflow_id} in database: {str(e)}")
        
        # Cache workflow in memory
        self.workflows[workflow_id] = workflow
        
        return workflow
    
    async def stop_workflow(self, workflow_id: str) -> bool:
        """
        Stop a running workflow.
        
        Args:
            workflow_id: ID of the workflow to stop
            
        Returns:
            True on success, False if workflow not found or already stopped
        """
        try:
            workflow = self.workflows.get(workflow_id)
            
            if not workflow:
                # Try to load the workflow
                workflow = await self.load_workflow(workflow_id)
                if not workflow:
                    return False
            
            # If workflow is not running or completed, abort
            if workflow.get("status") not in ["running", "completed"]:
                return False
            
            # Set status to stopped
            workflow["status"] = "stopped"
            workflow["last_activity"] = datetime.now().isoformat()
            
            self._add_log(workflow, "Workflow was manually stopped", "info", "workflow", "Workflow Management")
            
            # Save workflow
            self._save_workflow(workflow)
            
            return True
        except Exception as e:
            logger.error(f"Error stopping workflow {workflow_id}: {str(e)}")
            return False
    
    def _add_log(self, workflow: Dict[str, Any], message: str, log_type: str, agent_id: Optional[str] = None, agent_name: Optional[str] = None) -> None:
        """Add a log entry to the workflow."""
        # First, check if workflow is a string (ID) instead of dictionary
        if isinstance(workflow, str):
            # Try to load the workflow by ID
            workflow_id = workflow
            workflow = self.workflows.get(workflow_id)
            if not workflow:
                # Just log to the logger and return
                logger.info(f"Log (couldn't add to workflow {workflow_id}): {log_type} - {message}")
                return
        
        # Check if workflow is a dictionary
        if not isinstance(workflow, dict):
            logger.error(f"Invalid workflow type: {type(workflow)}. Expected dictionary.")
            # Just log to the logger and return
            logger.info(f"Log (couldn't add to workflow): {log_type} - {message}")
            return
        
        # Create log entry
        log_entry = {
            "id": f"log_{uuid.uuid4()}",
            "message": message,
            "type": log_type,
            "timestamp": datetime.now().isoformat(),
            "agent_id": agent_id,
            "agent_name": agent_name
        }
        
        # Add log entry to workflow
        if "logs" not in workflow:
            workflow["logs"] = []
        
        workflow["logs"].append(log_entry)
        
        # Update last activity
        workflow["last_activity"] = log_entry["timestamp"]
        
        # Save log entry to database if available
        if self.lucydom_interface:
            try:
                # Add workflow ID to log entry
                log_data = log_entry.copy()
                log_data["workflow_id"] = workflow["id"]
                
                self.lucydom_interface.create_workflow_log(log_data)
                logger.debug(f"Log entry for workflow {workflow['id']} saved to database")
            except Exception as e:
                logger.error(f"Error saving log entry for workflow {workflow['id']} to database: {str(e)}")
        
        # Also log to standard logger with the category prefix
        category_prefix = f"[{agent_name or agent_id or 'Workflow'}]" if agent_name or agent_id else ""
        log_message = f"{category_prefix} {message}"
        
        if log_type == "error":
            logger.error(log_message)
        elif log_type == "warning":
            logger.warning(log_message)
        else:
            logger.info(log_message)
    
    def get_workflow_status(self, workflow_id: str) -> Optional[Dict[str, Any]]:
        """
        Get the status of a workflow.
        
        Args:
            workflow_id: ID of the workflow
            
        Returns:
            Dictionary with status information or None if workflow not found
        """
        # Get from memory
        workflow = self.workflows.get(workflow_id)
        
        # If not in memory, load from database or file
        if not workflow:
            # Load from database if available
            if self.lucydom_interface:
                try:
                    workflow_data = self.lucydom_interface.get_workflow(workflow_id)
                    if workflow_data:
                        workflow = workflow_data
                except Exception as e:
                    logger.error(f"Error loading workflow status from database: {str(e)}")
            
            # If not in database, load from file
            if not workflow:
                try:
                    import json
                    workflow_path = os.path.join(self.results_dir, f"workflow_{workflow_id}.json")
                    if os.path.exists(workflow_path):
                        with open(workflow_path, 'r', encoding='utf-8') as f:
                            workflow = json.load(f)
                except Exception as e:
                    logger.error(f"Error loading workflow status from file: {str(e)}")
                    return None
        
        if not workflow:
            return None
        
        # Extract status information
        status_info = {
            "id": workflow.get("id"),
            "name": workflow.get("name", f"Workflow {workflow_id}"),
            "status": workflow.get("status"),
            "progress": 1.0 if workflow.get("status") in ["completed", "failed", "stopped"] else 0.5,
            "started_at": workflow.get("started_at"),
            "last_activity": workflow.get("last_activity"),
            "workflow_complete": workflow.get("status") == "completed",
            "current_round": workflow.get("current_round", 1),
            "data_stats": workflow.get("data_stats", {
                "total_processing_time": 0.0,
                "total_token_count": 0,
                "total_bytes_sent": 0,
                "total_bytes_received": 0
            })
        }
        
        return status_info
    
    def get_workflow_logs(self, workflow_id: str) -> Optional[List[Dict[str, Any]]]:
        """
        Get logs for a workflow.
        
        Args:
            workflow_id: ID of the workflow
            
        Returns:
            List of logs or None if workflow not found
        """
        # Get from memory
        workflow = self.workflows.get(workflow_id)
        
        # If not in memory, load from database
        if not workflow and self.lucydom_interface:
            try:
                logs = self.lucydom_interface.get_workflow_logs(workflow_id)
                return logs
            except Exception as e:
                logger.error(f"Error loading workflow logs from database: {str(e)}")
        
        # If not in database or no interface available, load from file
        if not workflow:
            try:
                import json
                workflow_path = os.path.join(self.results_dir, f"workflow_{workflow_id}.json")
                if os.path.exists(workflow_path):
                    with open(workflow_path, 'r', encoding='utf-8') as f:
                        workflow = json.load(f)
            except Exception as e:
                logger.error(f"Error loading workflow logs from file: {str(e)}")
                return None
        
        return workflow.get("logs", []) if workflow else None
    
    def get_workflow_messages(self, workflow_id: str) -> Optional[List[Dict[str, Any]]]:
        """
        Get messages for a workflow.
        
        Args:
            workflow_id: ID of the workflow
            
        Returns:
            List of messages or None if workflow not found
        """
        # Get from memory
        workflow = self.workflows.get(workflow_id)
        
        # If not in memory, load from database
        if not workflow and self.lucydom_interface:
            try:
                messages = self.lucydom_interface.get_workflow_messages(workflow_id)
                return messages
            except Exception as e:
                logger.error(f"Error loading workflow messages from database: {str(e)}")
        
        # If not in database or no interface available, load from file
        if not workflow:
            try:
                import json
                workflow_path = os.path.join(self.results_dir, f"workflow_{workflow_id}.json")
                if os.path.exists(workflow_path):
                    with open(workflow_path, 'r', encoding='utf-8') as f:
                        workflow = json.load(f)
            except Exception as e:
                logger.error(f"Error loading workflow messages from file: {str(e)}")
                return None
        
        return workflow.get("messages", []) if workflow else None

# Factory function for WorkflowManager
def get_workflow_manager(mandate_id: int = None, user_id: int = None, ai_service = None):
    """
    Get a WorkflowManager instance for the specified context.
    Reuses existing instances.
    
    Args:
        mandate_id: Mandate ID
        user_id: User ID
        ai_service: AI service
        
    Returns:
        WorkflowManager instance
    """
    from modules.lucydom_interface import get_lucydom_interface
    
    context_key = f"{mandate_id}_{user_id}"
    
    # LucyDOM interface for database access
    lucydom_interface = get_lucydom_interface(mandate_id, user_id)
    
    if context_key not in _workflow_managers:
        _workflow_managers[context_key] = WorkflowManager(
            mandate_id, 
            user_id, 
            ai_service, 
            lucydom_interface
        )
    
    # Update services if changed
    if ai_service is not None:
        _workflow_managers[context_key].ai_service = ai_service
    
    return _workflow_managers[context_key]

# Singleton factory for WorkflowManager instances per context
_workflow_managers = {}