import logging
from typing import Dict, Any, List
from modules.interfaces.interfaceAppModel import User
from modules.interfaces.interfaceChatModel import ChatWorkflow, UserInputRequest, TaskStep, TaskAction, ActionResult, ReviewResult, TaskPlan, WorkflowResult, TaskContext
from modules.chat.serviceCenter import ServiceCenter
from modules.interfaces.interfaceChatObjects import ChatObjects
from .handling.handlingTasks import HandlingTasks, WorkflowStoppedException

logger = logging.getLogger(__name__)

# ===== STATE MANAGEMENT AND VALIDATION CLASSES =====

class ChatManager:
    """Chat manager with improved AI integration and method handling"""
    
    def __init__(self, currentUser: User, chatInterface: ChatObjects):
        self.currentUser = currentUser
        self.chatInterface = chatInterface
        self.service: ServiceCenter = None
        self.workflow: ChatWorkflow = None
        self.handlingTasks: HandlingTasks = None

    async def initialize(self, workflow: ChatWorkflow) -> None:
        """Initialize chat manager with workflow"""
        self.workflow = workflow
        self.service = ServiceCenter(self.currentUser, self.workflow)
        self.handlingTasks = HandlingTasks(self.chatInterface, self.service, self.workflow)

    async def executeUnifiedWorkflow(self, userInput: UserInputRequest, workflow: ChatWorkflow) -> WorkflowResult:
        """Unified Workflow Execution"""
        try:
            logger.info(f"Starting unified workflow execution for workflow {workflow.id}")
            logger.debug(f"User request: {userInput.prompt}")
            
            # Phase 1: High-Level Task Planning
            logger.info("Phase 1: Generating task plan")
            task_plan = await self.handlingTasks.generateTaskPlan(userInput.prompt, workflow)
            if not task_plan or not task_plan.tasks:
                raise Exception("No tasks generated in task plan.")
            
            # Phase 2-5: For each task, execute and get results
            total_tasks = len(task_plan.tasks)
            logger.info(f"Phase 2: Executing {total_tasks} tasks")
            all_task_results = []
            previous_results = []
            for idx, task_step in enumerate(task_plan.tasks):
                # Pass task index to executeTask method
                current_task_index = idx + 1
                
                logger.info(f"Task {idx+1}/{total_tasks}: {task_step.objective}")
                
                # Create proper context object for this task
                task_context = TaskContext(
                    task_step=task_step,
                    workflow=workflow,
                    workflow_id=workflow.id,
                    available_documents=self.service.getAvailableDocuments(workflow),
                    available_connections=self.service.getConnectionReferenceList(),
                    previous_results=previous_results,
                    previous_handover=None,
                    improvements=[],
                    retry_count=0,
                    previous_action_results=[],
                    previous_review_result=None,
                    is_regeneration=False,
                    failure_patterns=[],
                    failed_actions=[],
                    successful_actions=[],
                    criteria_progress={
                        'met_criteria': set(),
                        'unmet_criteria': set(),
                        'attempt_history': []
                    }
                )
                
                # Execute task (this handles action generation, execution, and review internally)
                task_result = await self.handlingTasks.executeTask(task_step, workflow, task_context, current_task_index, total_tasks)
                # Handover
                handover_data = await self.handlingTasks.prepareTaskHandover(task_step, [], task_result, workflow)
                # Collect results
                all_task_results.append({
                    'task_step': task_step,
                    'task_result': task_result,
                    'handover_data': handover_data
                })
                # Update previous results for next task
                if task_result.success and task_result.feedback:
                    previous_results.append(task_result.feedback)
            
            # Final workflow result
            workflow_result = WorkflowResult(
                status="completed",
                completed_tasks=len(all_task_results),
                total_tasks=len(task_plan.tasks),
                execution_time=0.0,  # TODO: Calculate actual execution time
                final_results_count=len(all_task_results)
            )
            logger.info(f"Unified workflow execution completed successfully for workflow {workflow.id}")
            return workflow_result
        except WorkflowStoppedException:
            logger.info(f"Workflow {workflow.id} was stopped by user")
            return WorkflowResult(
                status="stopped", 
                completed_tasks=0,
                total_tasks=0,
                execution_time=0.0,
                final_results_count=0
            )
        except Exception as e:
            logger.error(f"Error in executeUnifiedWorkflow: {str(e)}")
            return WorkflowResult(
                status="failed", 
                completed_tasks=0,
                total_tasks=0,
                execution_time=0.0,
                final_results_count=0,
                error=str(e)
            )