import argparse
from typing import Dict, Any
from pydantic import AnyHttpUrl # type:ignore
from mcp.server.auth.settings import AuthSettings # type:ignore
from mcp.server.auth.provider import AccessToken, TokenVerifier # type:ignore
from mcp.server.fastmcp import Context, FastMCP # type:ignore
import sys
import time
import json

from config.logger import logger
from src.client_manager import MultiClientManager


# --- Inizializza il manager multi-cliente ---
try:
    client_manager = MultiClientManager()
    logger.info("🏢 Multi-Client Manager inizializzato")
except Exception as e:
    logger.error(f"❌ Errore inizializzazione manager: {e}")
    sys.exit(1)

class EnvironmentMultiClientTokenVerifier(TokenVerifier):
    "Token verifier che mappa token univoci ai client_id."

    def __init__(self):
        # Carica i token da variabili d'ambiente
        # Formato: MCP_TOKEN_<CLIENT_ID>=token_value
        self.token_to_client = {}
        
        # Scansiona tutte le variabili d'ambiente che iniziano con MCP_TOKEN_
        for key, value in os.environ.items():
            if key.startswith("MCP_TOKEN_"):
                client_id = key[10:].upper()  
                self.token_to_client[value] = client_id
        
        if not self.token_to_client:
            raise ValueError("Nessun token trovato! Definisci variabili MCP_TOKEN_<CLIENT_ID>")

    async def verify_token(self, token: str) -> AccessToken | None:
        "Verifica il token e restituisce il client_id associato."
        client_id = self.token_to_client.get(token)

        if client_id:
            access_token = AccessToken(
                token=token,
                client_id=client_id,
                scopes=["user"]
            )
            return access_token
        
        return None


# Initialize FastMCP server
mcp = FastMCP(
    "rag-server-multiclient-auth",
    # Token verifier for authentication
    token_verifier=EnvironmentMultiClientTokenVerifier(),
    # Auth settings for RFC 9728 Protected Resource Metadata
    auth=AuthSettings(
        issuer_url=AnyHttpUrl("https://auth.example.com"),  # Authorization Server URL
        resource_server_url=AnyHttpUrl("http://localhost:8000"),  # This server's URL
        required_scopes=["user"],
    ),
)

# --- MCP Tools ---
@mcp.tool()
def search_documents(query: str, limit: int) -> Dict[str, Any]:
    "Search documents for a specific client using RAG pipeline
    
    Args:
        query: The search query
        limit: Maximum number of results to return
    "

	# Get the client id from Beraer token or args params pass
	client_id = 'Client1'

    logger.info(f"🔍 Searching for client '{client_id}': '{query}' (limit: {limit})")
    start_time = time.time()
    
    try:
        results = client_manager.search_client_documents(
            client_id=client_id,
            query=query, 
            top_k=limit
        )
        
        search_time = time.time() - start_time
        
        return {
            "client_id": client_id,
            "query": query,
            "results": results,
            "total_found": len(results),
            "execution_time": f"{search_time:.2f}s",
            "status": "success"
        }
        
    except Exception as e:
        error_msg = f"Error during search for client {client_id}: {str(e)}"
        logger.error(error_msg)
        search_time = time.time() - start_time
        
        return {
            "client_id": client_id,
            "query": query,
            "results": [],
            "total_found": 0,
            "execution_time": f"{search_time:.2f}s",
            "status": "error",
            "error": error_msg
        }
    

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