import asyncio
import json
import os
import sys
from typing import Any, Literal

import asyncclick as click
from fastmcp import Client, FastMCP
from fastmcp.utilities.logging import configure_logging
from pydantic import BaseModel, Field

from fastmcp_agents.agent.basic import FastMCPAgent
from fastmcp_agents.cli.loader import get_config_for_bundled, get_config_from_file, get_config_from_url
from fastmcp_agents.cli.models import (
    AgentModel,
    AugmentedServerModel,
    OverriddenStdioMCPServer,
    ServerSettings,
)
from fastmcp_agents.errors.base import ContributionsWelcomeError
from fastmcp_agents.errors.cli import FastMCPAgentsError, NoConfigError
from fastmcp_agents.observability.logging import get_logger, setup_logging

logger = get_logger("cli")

if os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT") is not None:
    from fastmcp_agents.observability.otel import setup_otel

    setup_otel()


class PendingToolCall(BaseModel):
    name: str
    arguments: dict[str, Any]


class ToolCallResult(PendingToolCall):
    result: Any


class CliContext(BaseModel):
    server_settings: ServerSettings
    augmented_server_model: AugmentedServerModel = Field(default_factory=AugmentedServerModel)
    pending_tool_calls: list[PendingToolCall] = Field(default_factory=list)


@click.group()
@click.option(
    "--transport",
    type=click.Choice(["stdio", "sse", "streamable-http"]),
    default="stdio",
    help="The transport to use for the MCP server. (stdio, sse, streamable-http)",
)
@click.option(
    "--log-level",
    type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]),
    default="INFO",
    help="The logging level to use for the agent.",
)
@click.option("--agent-only", is_flag=True, help="Only run the agents, don't expose the tools to the client.")
@click.option("--tool-only", is_flag=True, help="Only run the tools, don't expose the agents to the client.")
@click.pass_context
def cli_base(
    ctx: click.Context,
    transport: Literal["stdio", "sse", "streamable-http"],
    log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
    agent_only: bool = False,
    tool_only: bool = False,
):
    """The base CLI for the FastMCP Agents. Configure logging, transport, and filtering."""

    # Setup FastMCP Agents Logging
    setup_logging(level=log_level)

    # Setup FastMCP Logging
    configure_logging(level=log_level, enable_rich_tracebacks=False)

    ctx.obj = CliContext(
        server_settings=ServerSettings(
            transport=transport,
            log_level=log_level,
            agent_only=agent_only,
            tool_only=tool_only,
        )
    )


@cli_base.group(name="cli", chain=True)
def cli_interface():
    """The group for the `CLI` part of the FastMCP Agents CLI.

    This allows building agents and wrapping servers entirely from the command line.
    """


@cli_interface.command(name="agent")
@click.option("--name", help="The name of the agent to wrap.")
@click.option("--description", help="The description of the agent to wrap.")
@click.option("--instructions", help="The instructions of the agent to wrap.")
@click.option("--allowed-tools", help="A comma separated list of the tools that the agent can use.")
@click.option("--blocked-tools", help="A comma separated list of the tools that the agent cannot use.")
@click.pass_context
def cli_build_agent(
    ctx: click.Context,
    name: str,
    description: str,
    instructions: str,
    allowed_tools: str,
    blocked_tools: str,
):
    """Build an agent over the command line."""
    split_allowed_tools = allowed_tools.split(",") if allowed_tools else None
    split_blocked_tools = blocked_tools.split(",") if blocked_tools else None

    cli_context: CliContext = ctx.obj

    cli_context.augmented_server_model.agents.append(
        AgentModel(
            name=name,
            description=description,
            default_instructions=instructions,
            allowed_tools=split_allowed_tools,
            blocked_tools=split_blocked_tools,
        )
    )


@cli_base.group(name="config", chain=True)
@click.option("--file", type=click.Path(exists=True), help="The config file to use.")
@click.option("--url", type=str, help="The URL of the config file to use.")
@click.option("--bundled", type=str, help="The bundled server to use.")
@click.pass_context
def cli_with_config(ctx: click.Context, file: str | None = None, url: str | None = None, bundled: str | None = None):
    """Load a config file from the command line."""

    cli_context: CliContext = ctx.obj

    if url:
        cli_context.augmented_server_model = get_config_from_url(url)
    elif file:
        cli_context.augmented_server_model = get_config_from_file(file)
    elif bundled:
        cli_context.augmented_server_model = get_config_for_bundled(bundled)
    else:
        raise NoConfigError


async def handle_pending_tool_calls(
    mcp_clients: list[Client], server: FastMCP, pending_tool_calls: list[PendingToolCall]
) -> list[ToolCallResult]:
    results: list[ToolCallResult] = []

    server_client = Client(transport=server)

    async with server_client:
        for pending_tool_call in pending_tool_calls:
            result = await server_client.call_tool(pending_tool_call.name, pending_tool_call.arguments)
            results.append(
                ToolCallResult(name=pending_tool_call.name, arguments=pending_tool_call.arguments, result=result),
            )

    for mcp_client in mcp_clients:
        await mcp_client.close()

    for result in results:
        logger.info(f"Tool {result.name} with arguments {result.arguments} returned result: {result.result}")

    return results


async def run_server_or_call_tools(
    agents: list[FastMCPAgent],
    mcp_clients: list[Client],
    server: FastMCP,
    pending_tool_calls: list[PendingToolCall],
    transport: Literal["stdio", "sse", "streamable-http"],
):
    if pending_tool_calls:
        await handle_pending_tool_calls(mcp_clients, server, pending_tool_calls)
    else:
        await server.run_async(transport=transport)

    total_token_usage = sum(agent.llm_link.token_usage for agent in agents)
    logger.info(f"Total token usage: {total_token_usage}")


@cli_with_config.command(name="run")
@click.pass_context
async def cli_with_config_run(ctx: click.Context):
    """Run the server."""
    cli_context: CliContext = ctx.obj

    agents, mcp_clients, server = await cli_context.augmented_server_model.to_fastmcp_server(server_settings=cli_context.server_settings)

    await run_server_or_call_tools(
        agents=agents,
        mcp_clients=mcp_clients,
        server=server,
        pending_tool_calls=cli_context.pending_tool_calls,
        transport=cli_context.server_settings.transport,
    )


@cli_interface.command(name="list")
@click.pass_context
async def list_tools(
    ctx: click.Context,
):
    """List the tools available on the server."""
    cli_context: CliContext = ctx.obj

    agents, mcp_clients, server = await cli_context.augmented_server_model.to_fastmcp_server(server_settings=cli_context.server_settings)

    tools = await server.get_tools()

    for mcp_client in mcp_clients:
        await mcp_client.close()

    logger.info("Listing tools:")
    for tool in tools.values():
        logger.info(f"Tool: {tool.name}")


@cli_interface.command(name="call", context_settings={"ignore_unknown_options": True, "allow_extra_args": True})
@click.argument("name", type=str)
@click.argument("parameters", type=str)
@click.pass_context
def call_tool(
    ctx: click.Context,
    name: str,
    parameters: str,
):
    """Add a tool call to the pending tool calls list."""
    cli_context: CliContext = ctx.obj

    arguments = json.loads(parameters)

    cli_context.pending_tool_calls.append(PendingToolCall(name=name, arguments=arguments))


@cli_base.command(name="shell")
@click.pass_context
def shell(ctx: click.Context):  # noqa: ARG001
    """Start a shell session with the server."""
    raise ContributionsWelcomeError(feature="shell")


@cli_interface.command(name="wrap", context_settings={"ignore_unknown_options": True, "allow_extra_args": True})
@click.option("--env", type=str, multiple=True, help="The environment variables to set for the server.")
@click.argument("direct-wrap-args", nargs=-1, type=click.UNPROCESSED)
@click.pass_context
async def wrap_server_run_mcp(
    ctx: click.Context,
    env: list[str],
    direct_wrap_args: list[str],
):
    """Take the last of the cli args and use them to run the mcp server and run it."""
    cli_context: CliContext = ctx.obj

    command = direct_wrap_args[0]
    args = direct_wrap_args[1:] if len(direct_wrap_args) > 1 else []
    environment = os.environ.copy()

    for env_var in env:
        key, value = env_var.split("=")
        environment[key] = value

    cli_context.augmented_server_model.mcpServers = {
        "main": OverriddenStdioMCPServer(
            command=command,
            args=args,
            env=environment,
        )
    }

    agents, mcp_clients, server = await cli_context.augmented_server_model.to_fastmcp_server(server_settings=cli_context.server_settings)

    await run_server_or_call_tools(
        agents=agents,
        mcp_clients=mcp_clients,
        server=server,
        pending_tool_calls=cli_context.pending_tool_calls,
        transport=cli_context.server_settings.transport,
    )


cli_with_config.add_command(call_tool)
cli_with_config.add_command(list_tools)


def run_mcp():
    try:
        asyncio.run(cli_base())
    except FastMCPAgentsError:
        logger.exception(msg="An error occurred while running FastMCPAgents")
        sys.exit(1)
    except KeyboardInterrupt:
        logger.info("Keyboard interrupt received, shutting down...")
        sys.exit(0)
    except Exception:
        logger.exception(msg="An unknown error occurred while running FastMCPAgents")
        sys.exit(1)


if __name__ == "__main__":
    run_mcp()
