from typing import List, Dict, Any, Optional, Type
import json
from mcp_agent.workflows.llm.augmented_llm_openai import (
    OpenAIAugmentedLLM,
    ChatCompletionMessageParam,
    ChatCompletionMessage,
)
from mcp_agent.workflows.llm.augmented_llm import (
    RequestParams,
    ModelT,
    MessageTypes,
)
from mcp_agent.logging.logger import get_logger
from mcp_agent.tracing.telemetry import get_tracer


class SimplePlanThenExecuteLLM(OpenAIAugmentedLLM):
    """
    Simplified Plan-then-Execute implementation.
    
    This approach:
    1. Analyzes the query and available tools
    2. Creates a simple execution plan
    3. Executes the plan step by step
    """
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Use a simple logger name to avoid any potential issues
        self._plan_logger = get_logger("plan_execute")
    
    async def generate(
        self,
        message: MessageTypes,
        request_params: RequestParams | None = None,
    ) -> List[ChatCompletionMessage]:
        """
        Plan-then-Execute generation process
        """
        tracer = get_tracer(self.context)
        with tracer.start_as_current_span(
            f"{self.__class__.__name__}.generate"
        ) as span:
            span.set_attribute("execution_mode", "plan_then_execute")
            
            params = self.get_request_params(request_params)
            
            # Step 1: Generate execution plan
            self._plan_logger.info("Creating execution plan...")
            plan = await self._generate_simple_plan(message, params)
            
            # Step 2: Execute plan
            self._plan_logger.info(f"Executing plan with {len(plan['steps'])} steps...")
            results = await self._execute_simple_plan(plan, params)
            
            # Step 3: Generate final response
            final_response = await self._generate_final_response(
                plan, results, message, params
            )
            
            return final_response
    
    async def _generate_simple_plan(
        self, 
        message: MessageTypes, 
        params: RequestParams
    ) -> Dict[str, Any]:
        """Generate a simple execution plan"""
        # Get available tools
        tools_response = await self.agent.list_tools()
        tools_list = [f"- {tool.name}: {tool.description}" for tool in tools_response.tools[:20]]  # Limit tools shown
        
        planning_prompt = f"""You are a strategic planner. Create a step-by-step plan to complete this task.

Task: {message}

Available tools (showing first 20):
{chr(10).join(tools_list)}

Create a JSON plan with this exact format:
{{
  "analysis": "Brief analysis of what needs to be done",
  "steps": [
    {{
      "step_number": 1,
      "tool_name": "tool_name_here",
      "tool_args": {{"arg1": "value1"}},
      "description": "What this step does"
    }}
  ]
}}

Important: Return ONLY the JSON plan, no other text."""

        # Use parent's generate method to get the plan
        responses = await super().generate(planning_prompt, params)
        
        # Extract JSON from response
        response_text = responses[0].content if responses else "{}"
        
        try:
            # Try to parse JSON from response
            plan = json.loads(response_text)
        except:
            # If parsing fails, create a simple default plan
            self._plan_logger.warning("Failed to parse plan, using default")
            plan = {
                "analysis": "Analyzing project structure",
                "steps": [
                    {
                        "step_number": 1,
                        "tool_name": "list_directory",
                        "tool_args": {"path": "."},
                        "description": "List current directory"
                    }
                ]
            }
        
        return plan
    
    async def _execute_simple_plan(
        self, 
        plan: Dict[str, Any], 
        params: RequestParams
    ) -> Dict[int, Any]:
        """Execute the plan step by step"""
        results = {}
        
        for step in plan.get("steps", []):
            step_num = step.get("step_number", 1)
            tool_name = step.get("tool_name", "")
            tool_args = step.get("tool_args", {})
            description = step.get("description", "")
            
            self._plan_logger.info(f"Executing step {step_num}: {description}")
            
            try:
                # Execute tool
                result = await self.agent.call_tool(
                    name=tool_name,
                    arguments=tool_args
                )
                results[step_num] = {"success": True, "result": result}
                
            except Exception as e:
                self._plan_logger.error(f"Step {step_num} failed: {e}")
                results[step_num] = {"success": False, "error": str(e)}
        
        return results
    
    async def _generate_final_response(
        self,
        plan: Dict[str, Any],
        results: Dict[int, Any],
        original_message: MessageTypes,
        params: RequestParams
    ) -> List[ChatCompletionMessage]:
        """Generate the final response based on plan execution results"""
        # Format execution results
        results_summary = []
        for step_num, result in results.items():
            if result.get("success"):
                results_summary.append(f"Step {step_num}: Success")
            else:
                results_summary.append(f"Step {step_num}: Failed - {result.get('error', 'Unknown error')}")
        
        final_prompt = f"""Based on the execution results, provide a response to the user.

Original request: {original_message}

Execution summary:
{chr(10).join(results_summary)}

Provide a helpful response that addresses the user's request."""

        # Use the parent's generate method for the final response
        return await super().generate(final_prompt, params) 