import asyncio
from typing import List, Dict, Any
from .mcp_client import MCPClient

class ToolRegistry:
    """Manages connections to all required MCP servers and their tools."""
    def __init__(self, server_urls: List[str]):
        self.servers: Dict[str, MCPClient] = {url: MCPClient(url) for url in server_urls}
        self.tools: Dict[str, Dict[str, Any]] = {}

    async def discover_tools(self):
        """Discovers all available tools from all connected MCP servers."""
        discovery_tasks = [client.list_tools() for client in self.servers.values()]
        results = await asyncio.gather(*discovery_tasks, return_exceptions=True)

        for i, client in enumerate(self.servers.values()):
            server_tools = results[i]
            if isinstance(server_tools, list):
                for tool in server_tools:
                    self.tools[tool["name"]] = {"client": client, "description": tool["description"]}

    async def execute(self, tool_name: str, params: Dict[str, Any]) -> Dict[str, Any]:
        """Finds the correct MCP server and executes the tool."""
        if tool_name not in self.tools:
            raise ValueError(f"Tool '{tool_name}' not found in registry.")
        tool_info = self.tools[tool_name]
        return await tool_info["client"].execute_tool(tool_name, params)

    async def close_all(self):
        """Closes all client connections."""
        await asyncio.gather(*(client.close() for client in self.servers.values()))

class AIAgent:
    """A mock AI agent that generates a plan based on a goal."""
    def __init__(self, model_name: str = "mock-planner"):
        self.model = model_name

    async def generate_plan(self, goal: str, available_tools: List[str]) -> List[Dict[str, Any]]:
        """
        Generates a step-by-step plan.
        In a real application, this would involve a call to a powerful LLM.
        Here, we use a hardcoded plan for demonstration.
        """
        await asyncio.sleep(1)  # Simulate LLM thinking time
        
        # This is a mock plan. A real LLM would generate this dynamically.
        plan = [
            {"step": 1, "thought": "I need a place to store the code. I'll use the `create_repo` tool.", "tool": "create_repo", "params": {"name": "my-awesome-blog", "private": False}},
            {"step": 2, "thought": "Now I need to scaffold a new Next.js application inside a secure sandbox.", "tool": "execute_shell", "params": {"command": "npx create-next-app@latest my-awesome-blog --yes"}},
            {"step": 3, "thought": "The project is created. I should commit the initial code.", "tool": "execute_shell", "params": {"command": "cd my-awesome-blog && git init && git add . && git commit -m 'Initial commit'"}},
            {"step": 4, "thought": "The plan is complete. I will report success.", "tool": "report_success", "params": {"message": "Blog project scaffolded successfully."}}
        ]
        return plan

class ForgeApp:
    """The main orchestrator for the Forge application."""
    def __init__(self, goal: str, mcp_server_urls: List[str]):
        self.goal = goal
        self.planner = AIAgent()
        self.tool_registry = ToolRegistry(server_urls=mcp_server_urls)

    async def run(self):
        """
        Runs the agent and yields status updates as a generator.
        """
        yield "🚀 **Starting Forge... Initializing systems.**"
        await self.tool_registry.discover_tools()
        yield f"✅ **Tool Discovery Complete.** Found {len(self.tool_registry.tools)} tools."

        available_tool_names = list(self.tool_registry.tools.keys())
        yield f"🧠 **Generating a plan for your goal:** '{self.goal}'"
        plan = await self.planner.generate_plan(self.goal, available_tool_names)
        yield "📝 **Plan Generated!** Starting execution..."

        for task in plan:
            yield f"\n**[Step {task['step']}]** 🤔 **Thought:** {task['thought']}"
            
            if task["tool"] == "report_success":
                yield f"🎉 **Final Result:** {task['params']['message']}"
                break

            try:
                yield f"🛠️ **Action:** Executing tool `{task['tool']}` with params: `{task['params']}`"
                result = await self.tool_registry.execute(task["tool"], task["params"])
                
                if result.get("status") == "error":
                    yield f"❌ **Error:** {result.get('result', 'Unknown error')}"
                    yield "🛑 **Execution Halted due to error.**"
                    break
                else:
                    yield f"✅ **Observation:** {result.get('result', 'Tool executed successfully.')}"

            except Exception as e:
                yield f"❌ **Critical Error executing step {task['step']}:** {e}"
                yield "🛑 **Execution Halted due to critical error.**"
                break
        
        await self.tool_registry.close_all()
        yield "\n🏁 **Forge execution finished.**"