"""Helpers compartilhados para tools do MCP Server."""

from __future__ import annotations

import asyncio
import logging
from typing import Any

from services.llm_service import get_llm_service
from services.usage_logger import _calc_cost, get_usage_logger, ms_since

logger = logging.getLogger("tecjustica-mcp.helpers")

# --- Limites de modelo ---
CHARS_POR_TOKEN = 4  # Estimativa: 1 token ≈ 4 chars
LIMITE_GEMINI_CHARS = 3_200_000  # ~800K tokens (margem segura para 1M do modelo)
BATCH_SIZE_CHARS = 800_000  # ~200K tokens por lote map-reduce

# --- Constantes compartilhadas ---
SYNC_URL = "https://tecjusticamcp3.vercel.app/login"

MSG_PROCESSO_NAO_ENCONTRADO = (
    "Processo nao encontrado. Verifique o numero (formato CNJ: NNNNNNN-DD.AAAA.J.TR.OOOO) "
    f"ou sincronize em {SYNC_URL}"
)
MSG_DOCUMENTO_NAO_ENCONTRADO = (
    "Documento nao encontrado. Use ls_documentos(numero_processo) para obter IDs validos."
)

LIMIAR_EVICTION = 10_000  # chars (~2500 tokens)
MAX_PROCESSOS = 50
MAX_DOCUMENTOS = 100
MAX_MOVIMENTACOES = 100

# --- Categorias de documentos para filtro inteligente ---
CATEGORIAS_DOCS = {
    "decisorio": ["sentenc", "acordao", "acord", "decisao", "decis", "despacho"],
    "peticoes": ["petic", "petio", "inicial", "contest", "replica", "recurso", "agravo"],
    "provas": ["laudo", "pericia", "certid", "documento", "comprova"],
    "criminal": ["denunc", "denncia", "inquer", "inqurito", "defesa"],
}

# --- Instrucoes de resumo por tipo ---
_INSTRUCOES_EVICTION: dict[str, str] = {
    "documento": (
        "Resuma este documento juridico mantendo: partes, datas, valores, "
        "decisoes, pedidos e fundamentacao. Seja conciso e preciso."
    ),
    "analise": (
        "Resuma esta analise mantendo: conclusoes, riscos, recomendacoes "
        "e pontos criticos. Seja conciso e preciso."
    ),
    "lista": (
        "Resuma esta lista mantendo os itens mais relevantes com suas "
        "informacoes essenciais. Agrupe por categoria se possivel."
    ),
}


def _get_current_user() -> tuple[str | None, str | None]:
    """Extrai user_id e email do token de acesso atual.

    Returns:
        (user_id, user_email) — ambos podem ser None se nao autenticado.
    """
    try:
        from fastmcp.server.dependencies import get_access_token

        token = get_access_token()
        if token and token.claims:
            return token.claims.get("sub"), token.claims.get("email")
    except Exception:
        pass
    return None, None


def _extrair_snippet(conteudo: str, termo: str, chars: int = 200) -> str:
    """Extrai snippet ao redor do termo encontrado."""
    termo_lower = termo.lower()
    conteudo_lower = conteudo.lower()

    pos = conteudo_lower.find(termo_lower)
    if pos == -1:
        return conteudo[:chars] + "..."

    inicio = max(0, pos - chars // 2)
    fim = min(len(conteudo), pos + len(termo) + chars // 2)

    snippet = conteudo[inicio:fim]
    if inicio > 0:
        snippet = "..." + snippet
    if fim < len(conteudo):
        snippet = snippet + "..."

    return snippet


def _montar_contexto_documentos(
    docs: list[dict[str, Any]], max_chars: int = LIMITE_GEMINI_CHARS
) -> str:
    """Monta contexto formatado para Gemini."""
    partes = []
    chars = 0

    for doc in docs:
        nome = doc.get("nome") or "Documento"
        tipo = doc.get("tipo_nome") or ""
        data = str(doc.get("data_juntada"))[:10] if doc.get("data_juntada") else ""
        conteudo = doc.get("conteudo") or ""

        parte = f"## {nome}\n**Tipo:** {tipo} | **Data:** {data}\n\n{conteudo}\n"

        if chars + len(parte) > max_chars:
            continue

        partes.append(parte)
        chars += len(parte)

    return "\n---\n\n".join(partes)


def paginar_resultado(
    items: list[Any],
    total: int,
    limite: int,
    offset: int,
) -> dict[str, Any]:
    """
    Formata resultado paginado com metadata consistente.

    Args:
        items: Itens da pagina atual
        total: Total de itens disponiveis
        limite: Limite de itens por pagina
        offset: Offset atual

    Returns:
        Dict com total, limite, offset, has_more, resultados
    """
    return {
        "total": total,
        "limite": limite,
        "offset": offset,
        "has_more": offset + limite < total,
        "resultados": items,
    }


async def _smart_response(
    conteudo: str,
    tipo: str = "documento",
    max_chars: int = LIMIAR_EVICTION,
    tool_name: str | None = None,
    user_id: str | None = None,
    user_email: str | None = None,
) -> dict[str, Any]:
    """
    Eviction inteligente: resume conteudo grande com LLM em vez de truncar.

    Args:
        conteudo: Texto a processar
        tipo: Tipo de conteudo ("documento", "analise", "lista")
        max_chars: Limiar para ativar resumo
        tool_name: Nome da tool (para logging)
        user_id: ID do usuario (para logging)
        user_email: Email do usuario (para logging)

    Returns:
        Dict com conteudo (original ou resumido), foi_resumido, tamanho_original,
        e opcionalmente tamanho_final e nota
    """
    tamanho_original = len(conteudo)

    if tamanho_original <= max_chars:
        return {
            "conteudo": conteudo,
            "foi_resumido": False,
            "tamanho_original": tamanho_original,
        }

    llm = get_llm_service()
    instrucao = _INSTRUCOES_EVICTION.get(tipo, _INSTRUCOES_EVICTION["documento"])

    resposta = await llm.resumir_texto(
        texto=conteudo,
        max_tokens_saida=4000,
        instrucao=instrucao,
        model="gemini-flash",
        tool_name=tool_name,
        user_id=user_id,
        user_email=user_email,
    )

    resumo = resposta.text
    return {
        "conteudo": resumo,
        "foi_resumido": True,
        "tamanho_original": tamanho_original,
        "tamanho_final": len(resumo),
        "nota": f"Resumido de {tamanho_original:,} para {len(resumo):,} chars. "
        "Para conteudo completo, use read_documento(document_id) com max_chars maior.",
    }


def erro_instrutivo(erro: str, sugestao: str) -> dict[str, Any]:
    """
    Formata erro com sugestao de acao para o agent.

    Args:
        erro: Descricao do erro
        sugestao: Sugestao de como resolver

    Returns:
        Dict com sucesso=False, erro e sugestao
    """
    return {"sucesso": False, "erro": erro, "dica": sugestao}


def _dividir_em_lotes(
    docs: list[dict[str, Any]], batch_chars: int = BATCH_SIZE_CHARS
) -> list[list[dict[str, Any]]]:
    """Divide documentos em lotes respeitando limite de chars."""
    lotes: list[list[dict[str, Any]]] = []
    lote_atual: list[dict[str, Any]] = []
    chars_lote = 0

    for doc in docs:
        conteudo = doc.get("conteudo") or ""
        doc_chars = len(conteudo) + 200  # overhead do header formatado
        if lote_atual and chars_lote + doc_chars > batch_chars:
            lotes.append(lote_atual)
            lote_atual = []
            chars_lote = 0
        lote_atual.append(doc)
        chars_lote += doc_chars

    if lote_atual:
        lotes.append(lote_atual)

    return lotes


async def analisar_map_reduce(
    docs: list[dict[str, Any]],
    pergunta: str,
    perspectiva: str | None = None,
    tipo_processo: str = "civel",
    batch_chars: int = BATCH_SIZE_CHARS,
    max_tokens_resumo: int = 4000,
    max_tokens_final: int = 4000,
    tool_name: str | None = None,
    user_id: str | None = None,
    user_email: str | None = None,
) -> tuple[str, dict[str, Any]]:
    """
    Map-reduce: processa docs em lotes com Gemini Flash.

    Returns:
        (resposta_final, metadata) onde metadata inclui lotes, tokens, estrategia
    """
    import time

    _start = time.perf_counter()
    llm = get_llm_service()

    lotes = _dividir_em_lotes(docs, batch_chars)
    total_lotes = len(lotes)

    if total_lotes <= 1:
        # Cabe tudo num unico contexto
        contexto = _montar_contexto_documentos(docs)
        prefixo = f"TIPO DE PROCESSO: {tipo_processo.upper()}\n"
        if perspectiva:
            prefixo += f"PERSPECTIVA DE ANALISE: {perspectiva}\n"
        contexto = prefixo + "\n" + contexto

        resposta = await llm.query(
            pergunta=pergunta,
            contexto=contexto,
            max_tokens=max_tokens_final,
            tool_name=tool_name,
            user_id=user_id,
            user_email=user_email,
        )
        tp = resposta.tokens_prompt or 0
        tc = resposta.tokens_completion or 0
        cost_in, cost_out, cost_total = _calc_cost("google/gemini-2.5-flash", tp, tc)
        return resposta.text, {
            "estrategia": "contexto_unico",
            "lotes": 1,
            "docs_total": len(docs),
            "contexto_chars": len(contexto),
            "tokens_prompt": tp,
            "tokens_completion": tc,
            "custo": {
                "tokens_prompt": tp,
                "tokens_completion": tc,
                "tokens_total": tp + tc,
                "custo_input_usd": float(cost_in),
                "custo_output_usd": float(cost_out),
                "custo_total_usd": float(cost_total),
            },
        }

    # --- FASE MAP: resumos parciais em paralelo ---
    logger.info("Map-reduce: %d lotes para %d docs", total_lotes, len(docs))

    async def _processar_lote(lote: list[dict[str, Any]], n: int) -> tuple[str, int, int]:
        contexto_lote = _montar_contexto_documentos(lote)
        prompt_map = (
            f"Voce esta analisando o LOTE {n}/{total_lotes} de documentos "
            f"de um processo judicial ({tipo_processo}).\n\n"
            "Resuma os pontos mais importantes de cada documento neste lote:\n"
            "- Partes, datas, valores, decisoes\n"
            "- Pedidos, fundamentacao, conclusoes\n"
            "- Prazos e alertas\n\n"
            "Seja conciso mas completo. Este resumo sera usado para uma analise consolidada."
        )
        resp = await llm.query(
            pergunta=prompt_map,
            contexto=contexto_lote,
            max_tokens=max_tokens_resumo,
            tool_name=f"{tool_name or 'analisar'}:map_{n}",
            user_id=user_id,
            user_email=user_email,
        )
        return resp.text, resp.tokens_prompt or 0, resp.tokens_completion or 0

    # Executa MAP em paralelo (tolerante a falhas parciais)
    tarefas = [_processar_lote(lote, i + 1) for i, lote in enumerate(lotes)]
    resultados = await asyncio.gather(*tarefas, return_exceptions=True)

    # Separar sucessos de falhas e agregar tokens
    resumos: list[tuple[int, str]] = []
    lotes_falhos: list[int] = []
    total_tokens_prompt = 0
    total_tokens_completion = 0
    for i, resultado in enumerate(resultados):
        if isinstance(resultado, BaseException):
            lotes_falhos.append(i + 1)
            logger.error("Map-reduce: lote %d/%d falhou: %s", i + 1, total_lotes, resultado)
        else:
            text, tp, tc = resultado
            resumos.append((i + 1, text))
            total_tokens_prompt += tp
            total_tokens_completion += tc

    if not resumos:
        raise RuntimeError(
            f"Map-reduce: todos os {total_lotes} lotes falharam. Ultimo erro: {resultados[-1]}"
        )

    if lotes_falhos:
        logger.warning(
            "Map-reduce: %d/%d lotes falharam: %s",
            len(lotes_falhos),
            total_lotes,
            lotes_falhos,
        )

    # --- FASE REDUCE: consolidacao ---
    resumos_concatenados = "\n\n---\n\n".join(
        f"## Resumo do Lote {num}/{total_lotes}\n{resumo}" for num, resumo in resumos
    )

    if lotes_falhos:
        resumos_concatenados += (
            f"\n\n⚠️ NOTA: {len(lotes_falhos)} de {total_lotes} lotes falharam "
            f"(lotes {lotes_falhos}). A analise abaixo e baseada nos lotes disponiveis."
        )

    prompt_reduce = pergunta
    if perspectiva:
        prompt_reduce += f"\n\nPERSPECTIVA DE ANALISE: {perspectiva}"

    contexto_reduce = (
        f"TIPO DE PROCESSO: {tipo_processo.upper()}\n\n"
        f"Voce recebeu resumos parciais de {total_lotes} lotes de documentos do processo.\n"
        "Com base nesses resumos, responda a pergunta do usuario.\n\n"
        f"RESUMOS PARCIAIS:\n\n{resumos_concatenados}"
    )

    resposta_final = await llm.query(
        pergunta=prompt_reduce,
        contexto=contexto_reduce,
        max_tokens=max_tokens_final,
        tool_name=f"{tool_name or 'analisar'}:reduce",
        user_id=user_id,
        user_email=user_email,
    )

    # Agregar tokens da fase REDUCE
    total_tokens_prompt += resposta_final.tokens_prompt or 0
    total_tokens_completion += resposta_final.tokens_completion or 0

    duration = ms_since(_start)
    total_chars = sum(len(doc.get("conteudo") or "") for doc in docs)

    # Calcular custo
    cost_in, cost_out, cost_total = _calc_cost(
        "google/gemini-2.5-flash", total_tokens_prompt, total_tokens_completion
    )

    metadata: dict[str, Any] = {
        "estrategia": "map_reduce",
        "lotes": total_lotes,
        "lotes_sucesso": len(resumos),
        "lotes_falhos": lotes_falhos,
        "docs_total": len(docs),
        "total_chars_originais": total_chars,
        "contexto_reduce_chars": len(contexto_reduce),
        "duration_ms": duration,
        "custo": {
            "tokens_prompt": total_tokens_prompt,
            "tokens_completion": total_tokens_completion,
            "tokens_total": total_tokens_prompt + total_tokens_completion,
            "custo_input_usd": float(cost_in),
            "custo_output_usd": float(cost_out),
            "custo_total_usd": float(cost_total),
        },
    }

    # Log consolidado
    try:
        get_usage_logger().log_llm_call(
            tool_name=tool_name or "analisar",
            call_type="map_reduce",
            model_id="google/gemini-2.5-flash",
            model_alias="gemini-flash",
            context_chars=total_chars,
            duration_ms=duration,
            user_id=user_id,
            user_email=user_email,
            metadata={"lotes": total_lotes, "docs": len(docs)},
        )
    except Exception:
        pass

    return resposta_final.text, metadata
