from mcp.server.fastmcp import FastMCP
from datetime import datetime
from llama_index.core import VectorStoreIndex
from llama_index.core import (
    StorageContext,
    load_index_from_storage,
)
from typing import Dict, Optional
import json
import os
import aiohttp  # Necesario para las peticiones HTTP asíncronas

port = int(os.getenv("PORT", 7860))
mcp = FastMCP("OnBase", port=port)







DOCUMENTS_BASE_PATH = "./"
SOURCES = {
    "oms": "oms/",
    # "other": "other/",
    # "other2": "other2/"
}

# Cargar índices
indices: Dict[str, VectorStoreIndex] = {}
for source, rel_path in SOURCES.items():
    full_path = os.path.join(DOCUMENTS_BASE_PATH, rel_path)
    if os.path.exists(full_path):
        storage_context = StorageContext.from_defaults(persist_dir=full_path)
        index = load_index_from_storage(storage_context, index_id="vector_index")
        indices[source] = index


# Recurso dinámico para búsqueda
@mcp.resource(
    uri="retriever://documentos/{fuente}",
    name="DocumentRetriever",
    description="Retrieve documents from differente regulations using semantic search.",
    mime_type="application/json",
    tags={"llm", "retrieval"}
)
def retrieve_docs(fuente: str, query: str, top_k: int = 3) -> dict:
    """
    Parameters:
        fuente: Source name (for example: oms).
        query: Search query.
        top_k: Number of results to return.
    """
    if fuente not in indices:
        return {"error": f"Fuente '{fuente}' no disponible. Opciones: {list(indices.keys())}"}
    
    retriever = indices[fuente].as_retriever(similarity_top_k=top_k)
    nodes = retriever.retrieve(query)
    
    results = [
        {
            "content": node.get_content(),
            "metadata": node.metadata,
            "score": node.score
        }
        for node in nodes
    ]
    
    return {"results": results}



    

@mcp.tool()
async def search_tavily(
    query: str,
    days: int = 7,
    max_results: int = 1,
    include_answer: bool = False
) -> dict:
    """Perform a web search using the Tavily API.
    
    Args:
        query: Search query string (required)
        days: Restrict search to last N days (default: 7)
        max_results: Maximum results to return (default: 1)
        include_answer: Include a direct answer only when requested by the user (default: False)
        
    Returns:
        dict: Search results from Tavily
    """
    # Obtener la API key de las variables de entorno
    tavily_api_key = os.environ.get('TAVILY_API_KEY')
    if not tavily_api_key:
        raise ValueError("TAVILY_API_KEY environment variable not set")
    
    headers = {
        "Authorization": f"Bearer {tavily_api_key}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "query": query,
        "search_depth": "basic",
        "max_results": max_results,
        "days": days if days else None,
        "include_answer": include_answer
    }
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(
                "https://api.tavily.com/search",
                headers=headers,
                json=payload
            ) as response:
                response.raise_for_status()
                result = await response.json()
                return result
                
    except Exception as e:
        return {
            "error": str(e),
            "status": "failed",
            "query": query
        }

if __name__ == "__main__":
    mcp.run("sse")