import logging
import os

from langchain_core.messages import ToolMessage
from langgraph.types import Command
from mcp.server.fastmcp import FastMCP
from typing import Union, Any, List, Dict
from uuid import uuid4


os.makedirs("logs", exist_ok=True)
logging.basicConfig(
    filename='logs/captain_server.log',                    # Log file name
    filemode='a',                          # 'a' for append, 'w' to overwrite
    level=logging.INFO,                    # Minimum level to log
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)


mcp = FastMCP("CaptainServer")


@mcp.tool()
async def TeamFinalChoice(
    guesses: List[str],
):
    """Use this tool if you think that some of the informations are missing.
    Args:
        guesses: a list of words (string) that you think are associated to the word told by the Boss Agent."""

    logger.info(f"Final choices made: {guesses}")

    artifact = {"guesses": guesses}

    return Command(
        update={
            "guesses": guesses,
            "messages": [ToolMessage(
                content="I made my final choices: " + ", ".join(guesses),
                name="TeamFinalChoice",
                artifact=artifact,
                tool_call_id=str(uuid4())
            )],
        }
    )


@mcp.tool()
async def Call_Agent_1(
    message: Union[str, Dict, List, Any],
):
    """Use this tool to call Agent 1 for help.
    Args:
        message: a message for Agent 1."""

    logger.info(f"Calling Agent 1 with message: {message}")

    return Command(
        update={
            "round_messages": [message],
            "messages": [ToolMessage(
                content=message,
                name="Call_Agent_1",
                tool_call_id=str(uuid4())
            )],
        }
    )


@mcp.tool()
async def Call_Agent_2(
    message: Union[str, Dict, List, Any],
):
    """Use this tool to call Agent 2 for help.
    Args:
        message: a message for Agent 2."""

    logger.info(f"Calling Agent 2 with message: {message}")

    return Command(
        update={
            "round_messages": [message],
            "messages": [ToolMessage(
                content=message,
                name="Call_Agent_2",
                tool_call_id=str(uuid4())
            )],
        }
    )


def start_captain_server():
    mcp.settings.port = 8001
    mcp.run(transport="stdio")


if __name__ == "__main__":
    start_captain_server()