import asyncio
import logging
import shutil
import sys
from contextlib import AsyncExitStack
from typing import Any

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.types import TextContent

logger = logging.getLogger(__name__)
logging.basicConfig(stream=sys.stderr, level=logging.INFO)


class LeanMCPClient:
    """Cliente programatico para o servidor MCP lean4-formal-verifier.

    Conecta ao servidor via stdio, envia codigo Lean para verificacao
    e retorna o resultado ([QED], [FEEDBACK DO COMPILADOR] ou [ERRO_SISTEMA]).
    """

    def __init__(self) -> None:
        self.session: ClientSession | None = None
        self._exit_stack = AsyncExitStack()

    async def connect(
        self,
        server_script: str | None = None,
        lean_project_path: str | None = None,
    ) -> list[str]:
        """Conecta ao servidor MCP lean4-formal-verifier.

        Args:
            server_script: Caminho para o script do servidor (.py).
                          Se None, usa o entry point 'lean-mcp-server'.
            lean_project_path: Caminho para o projeto Lean 4.
                              Passado como env var LEAN_PROJECT_PATH ao servidor.

        Returns:
            Lista de nomes das tools disponiveis no servidor.
        """
        if server_script:
            command = sys.executable
            args = [server_script]
        elif shutil.which("lean-mcp-server"):
            command = "lean-mcp-server"
            args = []
        else:
            command = sys.executable
            args = ["-m", "lean_mcp.server"]

        env: dict[str, str] | None = None
        if lean_project_path:
            env = {"LEAN_PROJECT_PATH": lean_project_path}

        server_params = StdioServerParameters(
            command=command,
            args=args,
            env=env,
        )

        stdio_transport = await self._exit_stack.enter_async_context(stdio_client(server_params))
        read_stream, write_stream = stdio_transport
        self.session = await self._exit_stack.enter_async_context(
            ClientSession(read_stream, write_stream)
        )

        await self.session.initialize()

        response = await self.session.list_tools()
        tool_names = [tool.name for tool in response.tools]
        logger.info("Conectado ao servidor com tools: %s", tool_names)
        return tool_names

    async def verify(self, lean_code: str, timeout: int = 120) -> str:
        """Envia codigo Lean para verificacao formal.

        Args:
            lean_code: Codigo Lean 4 completo (incluindo imports se necessario).
            timeout: Timeout em segundos para compilacao.

        Returns:
            String com o resultado: [QED], [FEEDBACK DO COMPILADOR] ou [ERRO_SISTEMA].

        Raises:
            RuntimeError: Se o cliente nao esta conectado.
        """
        if not self.session:
            raise RuntimeError("Cliente nao conectado. Chame connect() primeiro.")

        result = await self.session.call_tool(
            "verify_lean_theorem",
            arguments={"lean_code": lean_code, "timeout": timeout},
        )

        texts: list[str] = []
        for content in result.content:
            if isinstance(content, TextContent):
                texts.append(content.text)
            else:
                texts.append(str(content))

        return "\n".join(texts)

    async def list_tools(self) -> list[dict[str, Any]]:
        """Lista as tools disponiveis no servidor.

        Returns:
            Lista de dicionarios com name, description e inputSchema de cada tool.

        Raises:
            RuntimeError: Se o cliente nao esta conectado.
        """
        if not self.session:
            raise RuntimeError("Cliente nao conectado. Chame connect() primeiro.")

        response = await self.session.list_tools()
        return [
            {
                "name": tool.name,
                "description": tool.description,
                "input_schema": tool.inputSchema,
            }
            for tool in response.tools
        ]

    async def close(self) -> None:
        """Fecha a conexao com o servidor."""
        await self._exit_stack.aclose()

    async def __aenter__(self) -> "LeanMCPClient":
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.close()


async def _demo() -> None:
    """Exemplo de uso do cliente."""
    if len(sys.argv) < 2:
        print(
            "Uso: python -m lean_mcp.client <caminho_projeto_lean> [caminho_server.py]",
            file=sys.stderr,
        )
        sys.exit(1)

    lean_project = sys.argv[1]
    server_script = sys.argv[2] if len(sys.argv) > 2 else None

    async with LeanMCPClient() as client:
        tools = await client.connect(
            server_script=server_script,
            lean_project_path=lean_project,
        )
        print(f"Tools disponiveis: {tools}")

        code = "theorem modus_ponens (P Q : Prop) (hp : P) (hpq : P → Q) : Q := hpq hp"
        print(f"\nVerificando: {code}")
        result = await client.verify(code)
        print(f"Resultado: {result}")


if __name__ == "__main__":
    asyncio.run(_demo())
