from utils import get_mind_map, process_file, query_index
from fastmcp import FastMCP
from typing import List, Union, Literal
import os
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

mcp: FastMCP = FastMCP(name="MCP For NotebookLM")

# Check if all required environment variables are set
REQUIRED_ENV_VARS = [
    "LLAMACLOUD_API_KEY",
    "EXTRACT_AGENT_ID",
    "LLAMACLOUD_PIPELINE_ID",
    "OPENAI_API_KEY"
]

missing_vars = [var for var in REQUIRED_ENV_VARS if not os.getenv(var)]
if missing_vars:
    logger.error(f"Missing required environment variables: {', '.join(missing_vars)}")
    SERVICES_AVAILABLE = False
else:
    SERVICES_AVAILABLE = True


@mcp.tool(
    name="process_file_tool",
    description="This tool is useful to process files and produce summaries, question-answers and highlights.",
)
async def process_file_tool(
    filename: str,
) -> Union[str, Literal["Sorry, your file could not be processed."]]:
    if not SERVICES_AVAILABLE:
        logger.error("Services not available due to missing environment variables")
        return "Sorry, your file could not be processed."

    try:
        notebook_model, text = await process_file(filename=filename)
        if notebook_model is None:
            return "Sorry, your file could not be processed."
        if text is None:
            text = ""
        return notebook_model + "\n%separator%\n" + text
    except Exception as e:
        logger.error(f"Error processing file: {e}", exc_info=True)
        return "Sorry, your file could not be processed."


@mcp.tool(name="get_mind_map_tool", description="This tool is useful to get a mind map")
async def get_mind_map_tool(
    summary: str, highlights: List[str]
) -> Union[str, Literal["Sorry, mind map creation failed."]]:
    if not SERVICES_AVAILABLE:
        logger.error("Services not available due to missing environment variables")
        return "Sorry, mind map creation failed."

    try:
        mind_map_fl = await get_mind_map(summary=summary, highlights=highlights)
        if mind_map_fl is None:
            return "Sorry, mind map creation failed."
        return mind_map_fl
    except Exception as e:
        logger.error(f"Error creating mind map: {e}", exc_info=True)
        return "Sorry, mind map creation failed."


@mcp.tool(name="query_index_tool", description="Query a LlamaCloud index.")
async def query_index_tool(question: str) -> str:
    if not SERVICES_AVAILABLE:
        logger.error("Services not available due to missing environment variables")
        return "Sorry, I was unable to find an answer to your question."

    try:
        response = await query_index(question=question)
        if response is None:
            return "Sorry, I was unable to find an answer to your question."
        return response
    except Exception as e:
        logger.error(f"Error querying index: {e}", exc_info=True)
        return "Sorry, I was unable to find an answer to your question."


if __name__ == "__main__":
    if not SERVICES_AVAILABLE:
        logger.warning("Starting MCP server with limited functionality due to missing environment variables")
    mcp.run(transport="streamable-http")