"""mcp_creation_agent.py

This module implements the McpCreationAgent class, which serves as the central coordinator of the Alita pipeline.
It orchestrates the iterative CodeReAct loop to:
    - Analyze an input natural language task.
    - Detect capability gaps via the MCPBrainstorm module.
    - Retrieve external resources using ResearchAgent if necessary.
    - Generate a self-contained Python script with ScriptGenerator.
    - Set up an isolated runtime environment and execute the script via CodeRunner.
    - Register successfully executed scripts as new Model Context Protocols (MCPs) using MCPRegistry.
All steps are logged for reproducibility and debugging.
"""
import asyncio
import logging
from typing import Any, Dict, List

from calita.mcp_creation.code_runner_e2b import CodeRunnerSandbox
from calita.mcp_creation.mcp_brainstorm import MCPBrainstorm
from calita.mcp_creation.mcp_registry import MCPRegistry
from calita.mcp_creation.research_agent import ResearchAgent
from calita.mcp_creation.script_generator import ScriptGenerator
from calita.utils.utils import handle_error


class McpCreationAgent:
    def __init__(self, config: Dict[str, Any]) -> None:
        """
        Initialize the ManagerAgent with configuration settings.

        This constructor stores the configuration dictionary and instantiates
        the following submodules using the same configuration:
            - ResearchAgent for intelligent information retrieval using LangGraph and MCP tools.
            - MCPBrainstorm for analyzing tasks and generating tool specifications.
            - ScriptGenerator for producing self-contained executable scripts.
            - CodeRunner (which internally uses EnvironmentManager) for executing scripts.
            - MCPRegistry for registering and reusing successfully generated MCPs.
        Also initializes the maximum number of iterations for the iterative refinement loop.

        Args:
            config (Dict[str, Any]): Configuration dictionary loaded from config.yaml.
        """
        self.config: Dict[str, Any] = config
        self.research_agent: ResearchAgent = ResearchAgent(config)
        self.mcp_brainstorm: MCPBrainstorm = MCPBrainstorm(config)
        self.script_generator: ScriptGenerator = ScriptGenerator(config)

        self.code_runner = CodeRunnerSandbox(config)
        self.mcp_registry: MCPRegistry = MCPRegistry(config)
        # Set maximum iterations for the iterative CodeReAct loop; default value is 3.
        self.max_iterations: int = int(config.get("code_max_iterations", 3))
        logging.info("ManagerAgent initialized with max_iterations=%d", self.max_iterations)

    def analyze_task(self, task: str, context: str = "") -> Dict[str, Any]:
        """
        Analyze the given task using the MCPBrainstorm module to detect capability gaps.

        This method calls the brainstorming function using a given task description and
        additional context data (if any), and returns a specification dictionary.
        The dictionary typically contains:
            - "capability_gap": A boolean indicating if a tool (MCP) is needed.
            - "mcp_spec": A textual description/specification for the missing tool.
            - "task_description": The original task description.
            - "dependencies": A list of dependencies required (if any; default is an empty list).
            - "search_query": A suggested query for external resource lookup if a gap exists.

        Args:
            task (str): The natural language task description.
            context (str, optional): Additional context information from previous iterations. Defaults to "".

        Returns:
            Dict[str, Any]: The specification dictionary generated by MCPBrainstorm.
        """
        try:
            spec: Dict[str, Any] = self.mcp_brainstorm.brainstorm(task, context)
            # Ensure the specification contains the task description.
            if "task_description" not in spec:
                spec["task_description"] = task
            # Provide default dependencies if not specified.
            if "dependencies" not in spec:
                spec["dependencies"] = []
            # If a capability gap is indicated and no search query is provided, use the task as the query.
            if spec.get("capability_gap", False) and "search_query" not in spec:
                spec["search_query"] = task
            logging.info("Task analysis complete. Specification: %s", spec)
            return spec
        except Exception as e:
            handle_error(e)
            return {"capability_gap": False, "mcp_spec": None, "task_description": task, "dependencies": []}

    def generate(self, task: str, context: str = "") -> Dict[str, Any]:
        return asyncio.run(self.async_generate(task, context))

    async def async_generate(self, task: str, context: str="") -> Dict[str, Any]:
        """
        Orchestrate the full workflow for the given task through an iterative refinement loop.

        The method performs the following steps iteratively until successful script execution or
        until the maximum number of iterations is reached:
            1. Analyze the task via MCPBrainstorm to obtain a specification.
            2. If a capability gap exists, invoke ResearchAgent.search using a search query from the spec.
            3. Generate an executable script using ScriptGenerator with the obtained specification and resources.
            4. Set up the execution environment based on specification dependencies and a unique environment name.
            5. Execute the generated script via CodeRunner.
            6. On success, register the script as an MCP via MCPRegistry and return its output.
            7. If execution fails, incorporate the error feedback into the context for the next iteration.

        Args:
            task (str): The natural language task to be solved.

        Returns:
            str: The final output produced by the successful script execution, or an error message after all iterations.
        """
        iteration: int = 0
        context_feedback: str = "" if context is  None else context
        final_output = {}
        result = {}
        while iteration < self.max_iterations:
            logging.info("Iteration %d: Processing task: %s", iteration + 1, task)

            # Task analysis phase – generate specification using current context.
            spec: Dict[str, Any] = self.analyze_task(task, context_feedback)

            # If a capability gap is detected, perform an Aggregated search to gather resource context.
            resources: Dict[str, Any] = {}
            if spec.get("capability_gap", False):
                search_query: str = spec.get("search_query", task)
                logging.info("Capability gap detected. Initiating aggregated search with query: '%s'", search_query)
                resources = await self.research_agent.retrieve(search_query)
                logging.info("Aggregated search results obtained: %s", resources)
            else:
                logging.info("No external resource search required for this iteration.")

            # Generate the executable script based on the specification and gathered resources.
            generate_code: Dict[str, str] = self.script_generator.generate_script(spec, resources, context)
            script: str = generate_code.get('script', '')
            requirements: str = generate_code.get('requirements', '')
            function_name: str = generate_code.get('function_name', 'default_func')
            verify_code: str = generate_code.get('verify_code', '')

            # Use extracted requirements from script_generator if available, otherwise fall back to spec dependencies
            dependencies: List[str] = []
            if requirements:
                # Parse requirements string into list of dependencies
                dependencies: List[str] = [req.strip() for req in requirements.split('\n') if req.strip()]
                logging.info("Using extracted requirements from script: %s", dependencies)
            # else: #masked by zhangx , can't use
            #     # Fall back to dependencies from spec (for backward compatibility)
            #     dependencies: List[str] = spec.get("dependencies", [])
            #     logging.info("Using dependencies from spec: %s", dependencies)

            # Execute the generated script in the prepared environment.
            output, status = self.code_runner.run_script(function_name, script, verify_code, dependencies)
            logging.debug("Script execution status: %s; Output: %s", status, output)

            if status:
                # On successful execution, register the successful MCP and return the output.
                succeed = await self.mcp_registry.async_register_mcp_tool(function_name, script, requirements)
                final_output = {
                    "script": verify_code,
                    "output": output,
                }
                if succeed:
                    final_output["mcp_tool_name"] = function_name
                logging.info("Successful execution in iteration %d. MCP registered as '%s'.", iteration + 1, function_name)
                logging.info(f"========= McpCreationAgent Final Results Output =======\n{output}")
                result = {"result": final_output}
                return result
            else:
                # On failure, append the error message to the context and iterate.
                context_feedback += f"Iteration {iteration + 1} error: {output}\n"
                logging.warning("Iteration %d failed with error: %s", iteration + 1, output)
                iteration += 1

        # After maximum iterations, return a failure message with the last error encountered.
        final_failure: str = f"Failed to complete task after {self.max_iterations} iterations. Last error: {output}"
        logging.error(final_failure)
        result = {"error": final_failure}
        return result

if __name__ == "__main__":
    from calita.utils.utils import get_global_config
    from calita.utils.utils import setup_logging

    config = get_global_config("config.yaml")

    # Setup logging configuration
    setup_logging(config)


    agent = McpCreationAgent(config)
    result = agent.generate("Create a function to sort a list of numbers, sort [6,8,7,5]")
    print(result)