import logging
from contextlib import AsyncExitStack
from typing import Annotated, Any

import anyio
from fastmcp import Client
from langchain_core.messages import ToolMessage
from langchain_core.tools import InjectedToolCallId
from langchain_core.tools.base import ArgsSchema, BaseTool
from langgraph.types import Command
from mcp import McpError
from pydantic import BaseModel, Field, PrivateAttr

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("all.stratus.tools")


async def _close_mcp_client(exit_stack: AsyncExitStack, tool_name: str) -> None:
    """Ignore transport teardown failures after a tool call has already completed."""
    try:
        await exit_stack.aclose()
    except Exception as e:
        logger.warning("Ignoring %s while closing MCP client for %s", type(e).__name__, tool_name)


async def _connect_client(client: Client, exit_stack: AsyncExitStack) -> None:
    """Enter the client context and verify the session is actually connected.

    fastmcp's Client._connect can return without error even when the underlying
    SSE transport failed, leaving _session as None.  Detect this and raise early
    so callers can retry with a fresh client.
    """
    await exit_stack.enter_async_context(client)
    if not client.is_connected():
        raise RuntimeError("MCP client entered context but session is not connected")


async def _call_mcp_with_retry(
    tool: "BaseTool",
    mcp_tool_name: str,
    arguments: dict | None = None,
    *,
    session_id: str | None = None,
    max_retries: int = 1,
) -> str:
    """Call an MCP tool via the tool's client, retrying with a fresh client on connection failure."""
    from clients.stratus.stratus_utils.str_to_tool import get_client

    client = tool._client
    last_error: Exception | None = None

    for attempt in range(1 + max_retries):
        exit_stack = AsyncExitStack()
        try:
            await _connect_client(client, exit_stack)
            kwargs = {"arguments": arguments} if arguments else {}
            result = await client.call_tool(mcp_tool_name, **kwargs)
            return "\n".join([part.text for part in result])
        except (
            RuntimeError,
            ConnectionError,
            OSError,
            anyio.BrokenResourceError,
            anyio.ClosedResourceError,
            McpError,
        ) as e:
            last_error = e
            logger.warning(
                "MCP connection failed for %s (attempt %d/%d): %s",
                tool.name,
                attempt + 1,
                1 + max_retries,
                e,
            )
            # Replace the client with a fresh one for the retry
            client = get_client(session_id)
            tool._client = client
        finally:
            await _close_mcp_client(exit_stack, tool.name)

    logger.error("MCP tool %s failed after %d retries: %s", tool.name, max_retries, last_error)
    return f"Error: MCP tool call failed after {max_retries} retries: {last_error}"


class ExecKubectlCmdSafelyInput(BaseModel):
    command: str = Field(
        description="The command you want to execute in a CLI to manage a k8s cluster. "
        "It should start with 'kubectl'. Converts natural language to kubectl commands and executes them. "
        "Can be used to get/describe/edit Kubernetes deployments, services, and other Kubernetes components. "
        "Only takes one query at a time. Keep queries simple and straight-forward. "
        "This tool cannot handle complex mutli-step queries. "
        "Remember that most kubectl queries require a namespace name. "
    )
    tool_call_id: Annotated[str, InjectedToolCallId]


class ExecKubectlCmdSafely(BaseTool):
    name: str = "exec_kubectl_cmd_safely"
    description: str = "this is a tool used to safely execute kubectl commands."
    args_schema: ArgsSchema | None = ExecKubectlCmdSafelyInput

    _client: Client = PrivateAttr()
    _session_id: str = PrivateAttr()

    def __init__(self, client: Client, session_id: str, **kwargs: Any):
        super().__init__(**kwargs)
        self._client = client
        self._session_id = session_id

    @property
    def session_id(self) -> str:
        return self._session_id

    def _run(self):
        raise AssertionError(f"{self.name} is an async method, you are running it as a sync method!")
        pass

    async def _arun(
        self,
        command: str,
        tool_call_id: Annotated[str, InjectedToolCallId],
    ) -> Command:
        logger.debug(f"tool_call_id in {self.name}: {tool_call_id}")
        logger.debug(
            f'calling mcp exec_kubectl_cmd_safely from langchain exec_kubectl_cmd_safely, with command: "{command}"'
        )
        text_result = await _call_mcp_with_retry(
            self, "exec_kubectl_cmd_safely", {"cmd": command}, session_id=self._session_id
        )
        update: dict = {"messages": [ToolMessage(content=text_result, tool_call_id=tool_call_id)]}
        if "Command Rejected" not in text_result:
            update["executed_commands"] = [command]
        return Command(update=update)


kubectl_read_only_cmds = [
    "kubectl api-resources",
    "kubectl api-version",
    # read only if not interactive (interactive commands are prohibited)
    "kubectl attach",
    "kubectl auth can-i",
    "kubectl cluster-info",
    "kubectl describe",
    "kubectl diff",
    "kubectl events",
    "kubectl explain",
    "kubectl get",
    "kubectl logs",
    "kubectl options",
    "kubectl top",
    "kubectl version",
    "kubectl config view",
    "kubectl config current-context",
    "kubectl config get",
]


class ExecReadOnlyKubectlCmdInput(BaseModel):
    command: str = Field(
        description=f"The read-only kubectl command you want to execute in a CLI "
        'to manage a k8s cluster. It should start with "kubectl". '
        f"Available Read-only Commands: {kubectl_read_only_cmds}"
    )
    tool_call_id: Annotated[str, InjectedToolCallId]


class ExecReadOnlyKubectlCmd(BaseTool):
    name: str = "exec_read_only_kubectl_cmd"
    description: str = "this is a tool used to execute read-only kubectl commands."
    args_schema: ArgsSchema | None = ExecReadOnlyKubectlCmdInput

    _client: Client = PrivateAttr()

    def __init__(self, client: Client, **kwargs: Any):
        super().__init__(**kwargs)
        self._client = client

    def _run(self):
        raise AssertionError(f"{self.name} is an async method, you are running it as a sync method!")
        pass

    async def _arun(
        self,
        command: str,
        tool_call_id: Annotated[str, InjectedToolCallId],
    ) -> Command:
        logger.debug(f"tool_call_id in {self.name}: {tool_call_id}")
        is_read_only = False
        for c in kubectl_read_only_cmds:
            if command.startswith(c):
                is_read_only = True
                break
        if not is_read_only:
            logger.debug(
                f"Agent is trying to exec a non read-only command {command} with tool exec_read_only_kubectl_cmd"
            )
            text_result = (
                f"Your command {command} is not a read-only kubectl command. "
                f"Available Read-only Commands: {kubectl_read_only_cmds}."
            )
        elif command.startswith("kubectl logs -f"):
            logger.debug("agent calling interactive read-only command")
            text_result = f"Your command {command} is an _interactive_ read-only kubectl command. It is not supported!"
        else:
            logger.debug(
                f"calling mcp exec_kubectl_cmd_safely from "
                f'langchain exec_read_only_kubectl_cmd, with command: "{command}"'
            )
            text_result = await _call_mcp_with_retry(self, "exec_kubectl_cmd_safely", {"cmd": command})
        return Command(
            update={
                "messages": [
                    ToolMessage(content=text_result, tool_call_id=tool_call_id),
                ]
            }
        )


class RollbackCommandCmdInput(BaseModel):
    tool_call_id: Annotated[str, InjectedToolCallId]


class RollbackCommand(BaseTool):
    name: str = "rollback_command"
    description: str = (
        "Use this function to roll back the last kubectl command "
        'you successfully executed with the "exec_kubectl_cmd_safely" tool.'
    )
    args_schema: ArgsSchema | None = RollbackCommandCmdInput

    _client: Client = PrivateAttr()

    def __init__(self, client: Client, **kwargs: Any):
        super().__init__(**kwargs)
        self._client = client

    def _run(self):
        raise AssertionError(f"{self.name} is an async method, you are running it as a sync method!")
        pass

    async def _arun(
        self,
        tool_call_id: Annotated[str, InjectedToolCallId],
    ) -> Command:
        logger.debug(f"tool_call_id in {self.name}: {tool_call_id}")
        logger.debug("calling langchain rollback_command")
        text_result = await _call_mcp_with_retry(self, "rollback_command")
        return Command(
            update={
                "rollback_stack": str(text_result),
                "messages": [
                    ToolMessage(content=text_result, tool_call_id=tool_call_id),
                ],
            }
        )


class GetPreviousRollbackableCmdInput(BaseModel):
    tool_call_id: Annotated[str, InjectedToolCallId]


class GetPreviousRollbackableCmd(BaseTool):
    name: str = "get_previous_rollbackable_cmd"
    description: str = (
        "Use this function to get a list of commands you "
        "previously executed that could be roll-backed. "
        'When you call "rollback_command" tool multiple times, '
        "you will roll-back previous commands in the order "
        "of the returned list."
    )
    args_schema: ArgsSchema | None = GetPreviousRollbackableCmdInput

    _client: Client = PrivateAttr()

    def __init__(self, client: Client, **kwargs: Any):
        super().__init__(**kwargs)
        self._client = client

    def _run(self):
        raise AssertionError(f"{self.name} is an async method, you are running it as a sync method!")
        pass

    async def _arun(
        self,
        tool_call_id: Annotated[str, InjectedToolCallId],
    ) -> Command:
        logger.debug(f"tool_call_id in {self.name}: {tool_call_id}")
        logger.debug("calling langchain get_previous_rollbackable_cmd")
        text_result = await _call_mcp_with_retry(self, "get_previous_rollbackable_cmd")
        if not text_result:
            text_result = "There is no previous rollbackable command."
        return Command(
            update={
                "messages": [
                    ToolMessage(content=text_result, tool_call_id=tool_call_id),
                ]
            }
        )
