"""
Tool Router - LiteLLM Custom Logger / Guardrail
===============================================
Intercepta as chamadas do proxy do LiteLLM e faz retrieval das tools,
reduzindo drasticamente o consumo de tokens.

Como funciona:
1. Indexa cada tool (name + description)
2. No pre-call hook, calcula similaridade com a query do usuário
3. Seleciona top-K tools (+ always_include)
4. Só envia essas tools para o LLM

Uso como CALLBACK (roda em todas as chamadas):
  litellm_settings:
    callbacks:
      - tool_router_logger.tool_router_instance

Uso como GUARDRAIL (ativa por request):
  litellm_settings:
    guardrails:
      - guardrail_name: tool_router
        litellm_params:
          guardrail: tool_router_logger.ToolRouterGuardrail
          mode: pre_call
          default_on: false

Ou programaticamente:
  litellm.callbacks = [CustomToolRouter(...)]
"""

import asyncio
import hashlib
import json
import logging
import os
import re
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.caching import DualCache

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Embedding local leve (sem dependência externa pesada)
# ---------------------------------------------------------------------------

class FastEmbedder:
    """
    Usa o modelo de embedding do LiteLLM (qualquer provider via API)
    pra evitar depender de sentence-transformers local.
    
    Alternativa: Cache as embeddings no Valkey para reuse entre chamadas.
    """
    
    def __init__(self, embedding_model: str = "text-embedding-ada-002", cache: Optional[DualCache] = None):
        self.embedding_model = embedding_model
        self.cache = cache
        self._local_cache: Dict[str, List[float]] = {}
    
    async def embed(self, texts: List[str], litellm_kwargs: Optional[Dict] = None) -> List[List[float]]:
        """Gera embeddings em lote usando o LiteLLM."""
        import litellm
        
        to_embed = []
        indices = []
        results = [None] * len(texts)
        
        # Verifica cache primeiro (se disponível)
        for i, text in enumerate(texts):
            text_hash = hashlib.md5(text.encode()).hexdigest()
            cache_key = f"tool_embed:{text_hash}"
            
            if cache_key in self._local_cache:
                results[i] = self._local_cache[cache_key]
            elif self.cache:
                cached = await self.cache.async_get_cache(cache_key)
                if cached is not None:
                    val = cached
                    if isinstance(val, str):
                        val = json.loads(val)
                    self._local_cache[cache_key] = val
                    results[i] = val
                else:
                    to_embed.append(text)
                    indices.append(i)
            else:
                to_embed.append(text)
                indices.append(i)
        
        if not to_embed:
            return results  # type: ignore
        
        # Faz chamada de embedding
        kwargs = {"model": self.embedding_model, "input": to_embed}
        if litellm_kwargs:
            kwargs.update(litellm_kwargs)
        
        try:
            response = await litellm.aembedding(**kwargs)
            for idx, emb_data in zip(indices, response.data):
                embedding = emb_data["embedding"] if isinstance(emb_data, dict) else emb_data.embedding
                results[idx] = embedding
                
                # Salva no cache
                text = to_embed[indices.index(idx)]
                text_hash = hashlib.md5(text.encode()).hexdigest()
                cache_key = f"tool_embed:{text_hash}"
                self._local_cache[cache_key] = embedding
                if self.cache:
                    await self.cache.async_set_cache(cache_key, embedding)
        except Exception as e:
            logger.warning(f"Embedding failed, falling back to empty: {e}")
            for idx in indices:
                results[idx] = [0.0] * 1536  # fallback
        
        return results  # type: ignore


# ---------------------------------------------------------------------------
# Similaridade textual inteligente (fallback sem embedding externo)
# ---------------------------------------------------------------------------

import string
from collections import Counter

# Mapa de sinônimos e lematização para português + inglês (tools geralmente
# têm descrições em inglês, queries em português)
_SYNONYM_MAP = {
    # verbos para substantivos (português)
    "pesquise": "pesquisa", "pesquisar": "pesquisa", "buscar": "pesquisa", "busca": "pesquisa",
    "procure": "pesquisa", "procurar": "pesquisa", "encontre": "pesquisa",
    "leia": "leitura", "ler": "leitura", "abrir": "leitura", "abra": "leitura",
    "escreva": "escrita", "escrever": "escrita", "salvar": "escrita",
    "execute": "execução", "executar": "execução", "rodar": "execução", "rode": "execução",
    "liste": "listagem", "listar": "listagem", "mostre": "listagem",
    "envie": "envio", "enviar": "envio", "mande": "envio",
    "calcule": "cálculo", "calcular": "cálculo", "calcular": "cálculo",
    # português → inglês (tools descriptions são em inglês)
    "hora": "time", "horas": "time", "horário": "time", "horario": "time",
    "relógio": "time", "relogio": "time",
    "tempo": "weather time",  # ambos
    "clima": "weather", "previsão": "weather", "previsao": "weather",
    "chuva": "weather", "temperatura": "weather", "temperatures": "weather",
    "cidade": "city",
    "arquivo": "file", "arquivos": "file", "documento": "file",
    "pasta": "directory", "diretório": "directory", "diretorio": "directory",
    "email": "email", "e-mail": "email", "mensagem": "email",
    "web": "web", "internet": "web", "site": "web", "pesquisar": "search",
    "bash": "bash", "terminal": "bash", "comando": "bash command", "shell": "bash",
    "linha": "line",
    "navegador": "browser",
    "sistema": "system", "arquivos": "files",
}


# Stopwords comuns em português e inglês (palavras muito curtas/frequentes)
_STOPWORDS = {
    "a", "e", "o", "em", "de", "da", "do", "na", "no", "um", "uma",
    "é", "ser", "tem", "com", "para", "por", "pra", "os", "as",
    "ao", "aos", "das", "dos", "nas", "nos", "que", "se", "mas",
    "mais", "muito", "como", "vai", "vou", "me", "te", "se",
    "the", "a", "an", "in", "on", "at", "to", "for", "of",
    "is", "it", "be", "with", "and", "or", "from", "by",
    "um", "uns", "uma", "umas",
}


def _tokenize(text: str) -> list:
    """Tokeniza e normaliza, retornando lista (não set) pra preservar frequência."""
    text = text.lower().translate(str.maketrans('', '', string.punctuation))
    tokens = re.findall(r'[a-záéíóúãõâêîôûç]+', text)
    # Remove stopwords
    tokens = [t for t in tokens if t not in _STOPWORDS and len(t) > 1]
    # Aplica sinônimos
    expanded = []
    for t in tokens:
        expanded.append(t)
        if t in _SYNONYM_MAP:
            expanded.extend(_SYNONYM_MAP[t].split())
    return expanded


def _jaccard_similarity(a: set, b: set) -> float:
    if not a or not b:
        return 0.0
    return len(a & b) / len(a | b)


def _tf_keyword_score(query_tokens: list, tool_text: str) -> float:
    """
    Score baseado em quantas palavras-chave da query aparecem no texto da tool.
    Penaliza tools que não contêm NENHUMA palavra relevante.
    """
    tool_lower = tool_text.lower()
    matches = sum(1 for q in query_tokens if q in tool_lower)
    if matches == 0:
        return 0.0
    # Normaliza pelo tamanho da query pra favorecer matches parciais
    return matches / max(len(set(query_tokens)), 1)


def _combined_text_similarity(query: str, tool_text: str) -> float:
    """
    Combina Jaccard (set overlap) + TF-keyword score pra dar
    similaridade mais robusta sem precisar de embeddings externos.
    """
    query_tokens = _tokenize(query)
    tool_tokens = _tokenize(tool_text)
    
    query_set = set(query_tokens)
    tool_set = set(tool_tokens)
    
    # Jaccard entre sets
    jac = _jaccard_similarity(query_set, tool_set)
    
    # TF-score: quantas palavras da query estão no texto da tool
    tf = _tf_keyword_score(query_tokens, tool_text)
    
    # Identifica match exato no nome da tool (peso máximo)
    query_words = set(w for w in query_tokens if len(w) > 2)
    tool_name = tool_text.split()[0] if tool_text else ""
    name_match = 1.0 if tool_name in query_words else 0.0
    
    # Se o nome da tool apareceu na query, é quase certeza
    if name_match > 0:
        return 1.0
    
    # Peso maior pro TF quando Jaccard é baixo (captura relações indiretas)
    if jac == 0 and tf > 0:
        return tf * 0.4  # meio caminho, mais conservador
    
    # Só retorna algo relevante se passou de um mínimo
    raw = max(jac, tf * 0.6)
    return raw


# ---------------------------------------------------------------------------
# CustomLogger principal
# ---------------------------------------------------------------------------

DEFAULT_TOOL_NAMES = {"get_time", "help", "ping", "status"}

class CustomToolRouter(CustomLogger):
    """
    Router semântico de tools para LiteLLM Proxy.
    
    No pre-call hook, faz retrieval das tools mais relevantes
    para a query do usuário, reduzindo o prompt de tools enviado ao LLM.
    
    Args:
        embedding_model: Modelo de embedding (ex: 'text-embedding-ada-002')
        top_k: Número de tools a incluir (default: 3)
        min_similarity: Similaridade mínima para incluir tool (default: 0.15)
        always_include_tools: Lista de nomes de tools sempre enviadas
        similarity_mode: "embedding" (preferido) ou "jaccard" (fallback leve)
        cache_host: Host do Valkey/Redis (opcional)
        cache_port: Porta do Valkey/Redis (opcional)
        cache_password: Senha (opcional)
    """
    
    def __init__(
        self,
        embedding_model: str = "text-embedding-ada-002",
        top_k: int = 3,
        min_similarity: float = 0.15,
        always_include_tools: Optional[list] = None,
        similarity_mode: str = "embedding",
        cache_host: Optional[str] = None,
        cache_port: Optional[str] = None,
        cache_password: Optional[str] = None,
    ):
        super().__init__()
        self.top_k = top_k
        self.min_similarity = min_similarity
        self.always_include_tools = set(always_include_tools or DEFAULT_TOOL_NAMES)
        self.similarity_mode = similarity_mode
        self.embedding_model = embedding_model
        
        # Cache compartilhado (inicializado no async_pre_call_hook)
        self._dual_cache: Optional[DualCache] = None
        self._embedder: Optional[FastEmbedder] = None
        
        # Cache local de embeddings de tools
        self._tool_embeddings: Dict[str, List[float]] = {}
        self._tool_texts: Dict[str, str] = {}
        
        logger.info(
            f"CustomToolRouter initialized: "
            f"top_k={top_k}, min_similarity={min_similarity}, "
            f"mode={similarity_mode}, model={embedding_model}, "
            f"always_include={self.always_include_tools}"
        )
    
    def _extract_user_message(self, messages: List[Dict]) -> str:
        """Extrai a última mensagem do usuário."""
        for msg in reversed(messages):
            role = msg.get("role", "")
            content = msg.get("content", "")
            if role == "user":
                if isinstance(content, list):
                    # Pode ser multimodal - extrai texto
                    texts = []
                    for part in content:
                        if isinstance(part, dict) and part.get("type") == "text":
                            texts.append(part.get("text", ""))
                    return " ".join(texts)
                return str(content)
        return ""
    
    def _build_tool_text(self, tool: Dict) -> str:
        """Cria um texto combinado para a tool (nome + descrição + parâmetros)."""
        func = tool.get("function", tool)
        name = func.get("name", "")
        description = func.get("description", "")
        parameters = func.get("parameters", {})
        
        # Concatena informações relevantes
        parts = [name, description]
        
        # Adiciona nomes dos parâmetros
        if isinstance(parameters, dict):
            props = parameters.get("properties", {})
            if isinstance(props, dict):
                param_names = " ".join(props.keys())
                if param_names:
                    parts.append(param_names)
        
        return " ".join(parts).lower()
    
    async def _build_tool_index(self, tools: List[Dict]):
        """Indexa as tools: gera/cacheia embeddings."""
        texts = {}
        for tool in tools:
            func = tool.get("function", tool)
            name = func.get("name", "")
            text = self._build_tool_text(tool)
            texts[name] = text
            self._tool_texts[name] = text
        
        # Gera embeddings para tools ainda não indexadas
        if self.similarity_mode == "embedding" and self._embedder:
            new_tools = [name for name in texts if name not in self._tool_embeddings]
            if new_tools:
                tool_texts_list = [texts[name] for name in new_tools]
                try:
                    embeddings = await self._embedder.embed(tool_texts_list)
                    for name, emb in zip(new_tools, embeddings):
                        self._tool_embeddings[name] = emb
                except Exception as e:
                    logger.warning(f"Tool indexing failed, using jaccard: {e}")
                    self.similarity_mode = "jaccard"
    
    async def _compute_similarities(
        self, query: str, tool_names: List[str]
    ) -> List[Tuple[str, float]]:
        """Calcula similaridade entre query e cada tool."""
        if not query.strip():
            return [(name, 0.0) for name in tool_names]
        
        if self.similarity_mode == "embedding" and self._embedder:
            try:
                query_emb = (await self._embedder.embed([query]))[0]
                similarities = []
                for name in tool_names:
                    tool_emb = self._tool_embeddings.get(name)
                    if tool_emb:
                        sim = cosine_similarity([query_emb], [tool_emb])[0][0]
                        similarities.append((name, float(sim)))
                    else:
                        similarities.append((name, 0.0))
                return similarities
            except Exception as e:
                logger.warning(f"Embedding similarity failed: {e}")
                # fallback pra jaccard
        
        # Similaridade combinada (fallback inteligente)
        similarities = []
        for name in tool_names:
            tool_text = self._tool_texts.get(name, name)
            sim = _combined_text_similarity(query, tool_text)
            similarities.append((name, sim))
        
        return similarities
    
    async def async_pre_call_hook(
        self,
        user_api_key_dict: UserAPIKeyAuth,
        cache: DualCache,
        data: dict,
        call_type: str,
    ) -> Optional[Union[Exception, str, dict]]:
        """
        Hook principal: intercepta a requisição e filtra tools.
        
        Retorna:
            - dict modificado → segue com as tools filtradas
            - None → continua sem modificação
            - Exception → rejeita a request
        """
        # Só processa completion
        if call_type not in ("completion", "acompletion"):
            return None
        
        messages = data.get("messages", [])
        tools = data.get("tools", [])
        
        # Se não tem tools ou só 1, não precisa fazer nada
        if not tools or len(tools) <= 1:
            return None
        
        # Inicializa cache e embedder
        self._dual_cache = cache
        if self._embedder is None:
            self._embedder = FastEmbedder(
                embedding_model=self.embedding_model,
                cache=cache,
            )
        
        # Extrai query do usuário
        user_query = self._extract_user_message(messages)
        
        # Indexa as tools (se necessário)
        await self._build_tool_index(tools)
        
        # Nomes de todas as tools
        all_tool_names = []
        for tool in tools:
            func = tool.get("function", tool)
            all_tool_names.append(func.get("name", ""))
        
        # Calcula similaridades
        similarities = await self._compute_similarities(user_query, all_tool_names)
        
        # Filtra: sempre inclui + top-k por similaridade
        selected_names = set()
        
        # 1. Tools always-include
        for name in self.always_include_tools:
            if name in all_tool_names:
                selected_names.add(name)
        
        # 2. Top-k por similaridade (acima do threshold)
        sorted_tools = sorted(similarities, key=lambda x: x[1], reverse=True)
        for name, sim in sorted_tools:
            if len(selected_names) >= self.top_k:
                break
            if sim >= self.min_similarity:
                selected_names.add(name)
        
        # Se não selecionou nada, pelo menos a top-1
        if not selected_names and sorted_tools:
            selected_names.add(sorted_tools[0][0])
        
        # Monta novo tools array
        new_tools = [t for t in tools if (
            t.get("function", t).get("name", "") in selected_names
        )]
        
        # Log do que aconteceu
        logger.info(
            f"ToolRouter: tools {len(tools)}→{len(new_tools)} "
            f"(query: '{user_query[:60]}...' ) "
            f"selected: {sorted(selected_names)} "
            f"similarities: {dict(sorted_tools)}"
        )
        
        # Retorna dados modificados
        if len(new_tools) < len(tools):
            data["tools"] = new_tools
            # Se reduziu drasticamente, também podemos avisar o modelo
            if len(new_tools) <= 1 and messages:
                system_msg = (
                    "Você tem acesso apenas a tools selecionadas para esta conversa. "
                    "Se precisar de uma tool específica que não está disponível, "
                    "explique qual precisa e o usuário poderá liberá-la."
                )
                # Adiciona system message sutil, sem poluir
                existing_system = False
                for msg in messages:
                    if msg.get("role") == "system":
                        existing_system = True
                        if system_msg not in msg.get("content", ""):
                            msg["content"] = msg.get("content", "") + "\n\n" + system_msg
                        break
                if not existing_system:
                    messages.insert(0, {"role": "system", "content": system_msg})
                data["messages"] = messages
        
        return data
    
    async def async_log_success_event(self, *args, **kwargs):
        """Apenas para logging/monitoramento."""
        pass


# ---------------------------------------------------------------------------
# Factory para criação do logger via string path
# ---------------------------------------------------------------------------

# =============================================================================
# Instância pronta + factory + wrapper de guardrail
# =============================================================================
#
# Uso como CALLBACK (recomendado — roda em todas as chamadas):
#   litellm_settings:
#     callbacks: ["tool_router_logger.tool_router_instance"]
#
# Uso como GUARDRAIL (ativa por request com `guardrails: ["tool_router"]`):
#   litellm_settings:
#     guardrails:
#       - guardrail_name: tool_router
#         litellm_params:
#           guardrail: tool_router_logger.ToolRouterGuardrail
#           mode: pre_call
#           default_on: false
#           top_k: 3
#           min_similarity: 0.15
#           similarity_mode: jaccard
#
# Parâmetros configuráveis via env vars:
#   TOOL_ROUTER_EMBEDDING_MODEL, TOOL_ROUTER_TOP_K, TOOL_ROUTER_MIN_SIMILARITY,
#   TOOL_ROUTER_ALWAYS_INCLUDE (csv), TOOL_ROUTER_SIMILARITY_MODE
# =============================================================================

def _router_params_from_env() -> dict:
    """Lê os parâmetros do router a partir de env vars (com defaults sãos)."""
    return {
        "embedding_model": os.environ.get("TOOL_ROUTER_EMBEDDING_MODEL", "text-embedding-3-small"),
        "top_k": int(os.environ.get("TOOL_ROUTER_TOP_K", "3")),
        "min_similarity": float(os.environ.get("TOOL_ROUTER_MIN_SIMILARITY", "0.15")),
        "always_include_tools": set(
            filter(None, os.environ.get("TOOL_ROUTER_ALWAYS_INCLUDE", "get_time,help,list_tools").split(","))
        ),
        "similarity_mode": os.environ.get("TOOL_ROUTER_SIMILARITY_MODE", "jaccard"),
    }


# Instância global reutilizável — é o que o proxy carrega via:
#   callbacks: ["tool_router_logger.tool_router_instance"]
tool_router_instance: CustomToolRouter = CustomToolRouter(**_router_params_from_env())


def get_custom_tool_router(**kwargs) -> CustomToolRouter:
    """
    Factory function para uso programático.

    Se nenhum kwarg for passado, retorna a instância global
    (equivalente a `tool_router_instance`). Se kwargs forem passados,
    cria uma instância dedicada com aqueles parâmetros.
    """
    global tool_router_instance
    if kwargs:
        return CustomToolRouter(**kwargs)
    return tool_router_instance


class ToolRouterGuardrail(CustomGuardrail):
    """
    Wrapper para usar o ToolRouter como guardrail do LiteLLM Proxy.

    Guardrails permitem ativar o roteamento de tools por request:
      curl ... -d '{"model": "gpt-4o-mini", "guardrails": ["tool_router"], ...}'

    Ou globalmente com `default_on: true` no config.yaml.
    """

    _ROUTER_PARAMS = (
        "embedding_model", "top_k", "min_similarity",
        "always_include_tools", "similarity_mode",
        "cache_host", "cache_port", "cache_password",
    )

    def __init__(self, guardrail_name: str = "tool_router", **kwargs):
        # `guardrail_name`, `event_hook`, `default_on` e demais campos do
        # guardrail são repassados à base; os parâmetros do router são
        # extraídos para criar a instância interna.
        super().__init__(guardrail_name=guardrail_name, **kwargs)
        router_params = {k: v for k, v in kwargs.items() if k in self._ROUTER_PARAMS}
        self._router = CustomToolRouter(**router_params)

    async def async_pre_call_hook(
        self,
        user_api_key_dict: UserAPIKeyAuth,
        cache: DualCache,
        data: dict,
        call_type: str,
    ) -> Optional[Union[Exception, str, dict]]:
        """Delega o filtro de tools ao CustomToolRouter."""
        return await self._router.async_pre_call_hook(
            user_api_key_dict=user_api_key_dict,
            cache=cache,
            data=data,
            call_type=call_type,
        )

    async def async_log_success_event(self, *args, **kwargs):
        pass