from typing import Any
from pydantic import BaseModel, Field
from .base import BaseTool
from core.agent import AgentManager


class PlanEnterArgs(BaseModel):
    """Arguments for entering plan mode"""

    plan: str = Field(
        ...,
        description="The plan you came up with, that you want to run by the user for approval",
    )


class PlanExitArgs(BaseModel):
    """Arguments for exiting plan mode"""

    pass


class PlanEnterTool(BaseTool):
    """Tool for entering plan mode"""

    name = "plan_enter"
    description = """Use this tool when you are in build mode and want to enter plan mode to create a detailed plan before implementation.

IMPORTANT: Only use this tool when the task requires planning the implementation steps of a task that requires writing code. For research tasks where you're gathering information, searching files, or understanding the codebase - do NOT use this tool.

Before using this tool, ensure your plan is clear and unambiguous. If there are multiple valid approaches or unclear requirements, ask the user for clarification first."""
    args_schema = PlanEnterArgs

    def __init__(self, agent_manager: AgentManager):
        self.agent_manager = agent_manager

    async def run(self, plan: str) -> str:
        """Execute the tool"""
        current = self.agent_manager.get_current_agent()
        if current and current.name == "plan":
            return "Already in plan mode."

        # Switch to plan mode
        self.agent_manager.switch_agent("plan")

        return f"Entering plan mode. Your plan:\n\n{plan}\n\nYou are now in readonly mode. You can only read files and search code."


class PlanExitTool(BaseTool):
    """Tool for exiting plan mode"""

    name = "plan_exit"
    description = """Use this tool when you are in plan mode and have finished presenting your plan and are ready to code.

This will switch you back to code/build mode where you can make file changes and run commands."""
    args_schema = PlanExitArgs

    def __init__(self, agent_manager: AgentManager):
        self.agent_manager = agent_manager

    async def run(self) -> str:
        """Execute the tool"""
        current = self.agent_manager.get_current_agent()
        if not current or current.name != "plan":
            return "Not in plan mode."

        # Switch to build mode
        target_agent = "build"

        self.agent_manager.switch_agent(target_agent)

        # Get switch message
        message = self.agent_manager.get_switch_message("plan", target_agent)

        return f"Exiting plan mode. Switched to {target_agent} mode.\n{message}"
