r"""
As Tools que falam com o CNJ.

São as mesmas duas da Unidade 1 — uma resolve o nome do assunto, a outra
busca os processos — só que agora o dado vem da API do DataJud e não de um
SQLite local.

Rode com:
    python exemplos\tools_cnj.py
"""

import json
import sys
import time

from cnj import consultar

sys.stdout.reconfigure(encoding="utf-8")

# Teto de resultados. Mesma ideia da Unidade 1: o modelo pede o que quiser,
# quem decide o máximo é o seu código.
LIMITE_MAXIMO = 20

# Nível de sigilo aceito nesta sessão.
#   0 = Público   1 = Segredo de Justiça   2 = Sigilo Absoluto
# A API Pública só publica processos de nível 0. O filtro abaixo não existe
# porque desconfiamos do CNJ: existe porque um filtro de acesso pertence ao
# SEU código, e não à sorte de a origem estar sempre certa.
TETO_SIGILO = 0

# Teto de movimentos devolvidos por processo. Existe pelo mesmo motivo que
# LIMITE_MAXIMO, e com urgência maior: um processo do índice do TJRJ chega a
# 613 movimentos, e devolver isso inteiro estoura a janela do modelo num passo.
LIMITE_MOVIMENTOS = 20

# O índice grava dataAjuizamento em DOIS formatos: ISO ("2025-03-06T13:16:00.000Z")
# e compacto ("19980417000000"). O Elasticsearch lê o compacto como milissegundos
# desde 1970, o que joga toda data compacta para o século 27 — e portanto acima de
# QUALQUER data ISO na ordenação. Medido no índice do TJRJ: o processo compacto mais
# antigo é de 1998 e ainda assim ordena acima do processo ISO mais novo, de 2025.
#
# Este corte separa as duas famílias. Nenhuma data real cai entre elas: as ISO ficam
# todas abaixo de 2100, e as compactas todas em 26xx.
CORTE_DE_FORMATO = "2100-01-01"

_CACHE_ASSUNTOS: list[str] = []


def _catalogo_de_assuntos() -> list[str]:
    """Todos os nomes de assunto que aparecem em processos do TJRJ."""
    if _CACHE_ASSUNTOS:
        return _CACHE_ASSUNTOS

    resposta = consultar(
        {
            "size": 0,
            "aggs": {"nomes": {"terms": {"field": "assuntos.nome.keyword", "size": 3000}}},
        },
        timeout=120,
    )
    baldes = resposta["aggregations"]["nomes"]["buckets"]
    _CACHE_ASSUNTOS.extend(b["key"] for b in baldes)
    return _CACHE_ASSUNTOS


def formatar_numero(numero: str) -> str:
    """Põe a máscara CNJ em um número de 20 dígitos.

    A API devolve 00290982020268190000; um juiz lê 0029098-20.2026.8.19.0000.
    """
    n = "".join(c for c in numero if c.isdigit())
    if len(n) != 20:
        return numero
    return f"{n[:7]}-{n[7:9]}.{n[9:13]}.{n[13]}.{n[14:16]}.{n[16:]}"


def chave_data(bruta: str) -> str:
    """Reduz qualquer um dos dois formatos do índice a AAAAMMDD.

    20260529185822           -> 20260529
    2025-03-06T13:16:00.000Z -> 20250306

    Nessa forma, comparar como texto é comparar como data — que é o que
    permite ordenar processos das duas famílias na mesma lista.
    """
    digitos = "".join(c for c in str(bruta or "") if c.isdigit())
    return digitos[:8]


def formatar_data(bruta: str) -> str:
    """20260529185822 -> 29/05/2026. Aceita também o formato ISO."""
    d = chave_data(bruta)
    if len(d) != 8:
        return bruta
    return f"{d[6:8]}/{d[4:6]}/{d[:4]}"


def listar_assuntos(termo: str) -> list:
    """Nomes de assunto do TJRJ que contêm o termo procurado.

    Devolve uma LISTA de strings — não um JSON serializado. Num Code Agent,
    o modelo escreve Python de verdade sobre o que a Tool devolve: se isto
    fosse uma string, `for assunto in listar_assuntos(...)` percorreria
    caracteres, e cada caractere viraria uma consulta ao CNJ.
    """
    alvo = termo.lower().strip()
    achados = [nome for nome in _catalogo_de_assuntos() if alvo in nome.lower()]

    if not achados:
        raise ValueError(
            f"Nenhum assunto contém '{termo}'. Tente um termo mais curto, "
            "por exemplo apenas 'violência' ou 'doméstica'."
        )

    return achados[:20]


def buscar_processos(assunto: str, quantidade: int = 5) -> list:
    """Processos mais recentes de um assunto, do mais novo para o mais antigo.

    O `assunto` precisa ser o nome EXATO, tal como devolvido por listar_assuntos.

    A data de ajuizamento vem do índice em dois formatos, e por isso a busca é
    feita em duas consultas e reordenada aqui — ver CORTE_DE_FORMATO.

    Devolve uma LISTA de dicionários, e não um dicionário com a lista dentro.
    A diferença importa porque um agente compõe resultados: chamado uma vez
    por assunto, ele vai juntar as respostas com `.extend()`. Listas se juntam;
    dicionários, quando percorridos, entregam as CHAVES. Cada processo carrega
    o próprio `assunto` justamente para sobreviver a essa junção.
    """
    quantidade = max(1, min(int(quantidade), LIMITE_MAXIMO))

    # Rejeita aqui o que o índice não conhece, antes de gastar uma chamada de
    # rede. Um assunto inválido é erro de quem chamou, e o erro deve custar
    # uma exceção local — não uma requisição ao CNJ.
    if assunto not in _catalogo_de_assuntos():
        raise ValueError(
            f"'{assunto}' não é um nome de assunto do TJRJ. "
            "Use listar_assuntos() para obter os nomes exatos."
        )

    # Uma consulta só, ordenada por dataAjuizamento, devolveria apenas a família
    # compacta: ela ocupa TODAS as primeiras posições, por milhares de documentos.
    # Então perguntamos duas vezes, uma por família de formato, e juntamos aqui.
    # Custa uma requisição a mais, e é a diferença entre "os mais recentes" e
    # "os que o índice por acaso põe na frente".
    achados = []
    for familia in (
        {"range": {"dataAjuizamento": {"lt": CORTE_DE_FORMATO}}},
        {"range": {"dataAjuizamento": {"gte": CORTE_DE_FORMATO}}},
    ):
        corpo = {
            "size": quantidade,
            "query": {
                "bool": {
                    "filter": [
                        # term em .keyword: comparação exata, sem quebrar em palavras.
                        {"term": {"assuntos.nome.keyword": assunto}},
                        {"range": {"nivelSigilo": {"lte": TETO_SIGILO}}},
                        familia,
                    ]
                }
            },
            "sort": [{"dataAjuizamento": {"order": "desc"}}],
            "_source": [
                "numeroProcesso",
                "dataAjuizamento",
                "nivelSigilo",
                "grau",
                "classe.nome",
                "orgaoJulgador.nome",
            ],
        }
        achados.extend(consultar(corpo)["hits"]["hits"])

    if not achados:
        return []

    # Reordena pela data normalizada. É este passo que intercala as duas
    # famílias — sem ele, os dois blocos ficariam simplesmente concatenados.
    achados.sort(
        key=lambda h: chave_data(h["_source"].get("dataAjuizamento")),
        reverse=True,
    )

    processos = []
    for h in achados:
        p = h["_source"]
        # Cinto e suspensório: se algo com sigilo escapar da origem, para aqui.
        if p.get("nivelSigilo", 99) > TETO_SIGILO:
            continue
        processos.append(
            {
                "assunto": assunto,
                "numero_processo": formatar_numero(p["numeroProcesso"]),
                "data_ajuizamento": formatar_data(p["dataAjuizamento"]),
                "grau": p.get("grau"),
                "classe": p.get("classe", {}).get("nome"),
                "orgao_julgador": p.get("orgaoJulgador", {}).get("nome"),
            }
        )

    # Duas consultas trouxeram até 2x o pedido; o corte final é aqui.
    return processos[:quantidade]


def ultimos_movimentos(numero_processo: str, quantidade: int = 5) -> list:
    """Os últimos andamentos de um processo, do mais recente para o mais antigo.

    Aceita o número com ou sem máscara: 0018683-81.2020.8.19.0066 e
    00186838120208190066 dão no mesmo.

    Um número de processo não identifica UM documento no DataJud: o mesmo
    processo aparece uma vez por grau em que tramitou, cada um com a sua
    classe e a sua lista de movimentos. Esta Tool junta os graus e ordena
    tudo por data, porque é assim que a pergunta é feita — "o que andou
    neste processo" não costuma ser uma pergunta sobre instância.

    Devolve uma LISTA de dicionários, já cortada em `quantidade`. O corte é o
    ponto da função: a rede traz o processo inteiro, e é aqui, antes da
    Observation, que ele deixa de ser grande.
    """
    quantidade = max(1, min(int(quantidade), LIMITE_MOVIMENTOS))

    # Guarda antes da rede, de novo: um número que não tem 20 dígitos não é
    # número de processo, e não merece uma requisição ao CNJ para descobrir isso.
    digitos = "".join(c for c in str(numero_processo) if c.isdigit())
    if len(digitos) != 20:
        raise ValueError(
            f"'{numero_processo}' não é um número de processo. "
            "São 20 dígitos, com ou sem máscara."
        )

    resposta = consultar(
        {
            # size 10 e para os graus, não para os movimentos: um processo
            # tramita em poucas instâncias, mas nunca em dez.
            "size": 10,
            "query": {
                "bool": {
                    "filter": [
                        {"term": {"numeroProcesso": digitos}},
                        {"range": {"nivelSigilo": {"lte": TETO_SIGILO}}},
                    ]
                }
            },
            # Duas reduções, em lugares diferentes. Esta é a primeira e a mais
            # barata: `complementosTabelados` é quase metade do peso de
            # movimentos[] e não entra na resposta a "o que andou no processo".
            # Cortado aqui, nem chega a virar tráfego de rede.
            "_source": {
                "includes": ["numeroProcesso", "grau", "nivelSigilo", "classe.nome", "movimentos"],
                "excludes": ["movimentos.complementosTabelados"],
            },
        }
    )
    achados = resposta["hits"]["hits"]

    if not achados:
        raise ValueError(
            f"Nenhum processo público com o número {formatar_numero(digitos)} "
            "no índice do TJRJ."
        )

    movimentos = []
    for h in achados:
        p = h["_source"]
        if p.get("nivelSigilo", 99) > TETO_SIGILO:
            continue
        for m in p.get("movimentos") or []:
            movimentos.append(
                {
                    "numero_processo": formatar_numero(p["numeroProcesso"]),
                    "grau": p.get("grau"),
                    "classe": p.get("classe", {}).get("nome"),
                    "codigo": m.get("codigo"),
                    "movimento": m.get("nome"),
                    # dataHora vem ISO: 2021-06-11T00:00:00.000Z. formatar_data
                    # aceita os dois formatos do índice, então não há conversão aqui.
                    "data": formatar_data(m.get("dataHora")),
                    # Guardado só para ordenar; o formato ISO da API ordena
                    # certo como texto, o formato brasileiro não ordenaria.
                    "_ordem": m.get("dataHora") or "",
                }
            )

    movimentos.sort(key=lambda m: m["_ordem"], reverse=True)

    for m in movimentos:
        del m["_ordem"]

    return movimentos[:quantidade]


if __name__ == "__main__":
    inicio = time.time()
    print("Catálogo de assuntos:", len(_catalogo_de_assuntos()), "nomes distintos")
    print(f"(a primeira chamada levou {time.time() - inicio:.1f}s)\n")

    print("listar_assuntos('violência doméstica'):")
    print(listar_assuntos("violência doméstica"), "\n")

    print("buscar_processos('Violência Doméstica Contra a Mulher', 3):")
    for processo in buscar_processos("Violência Doméstica Contra a Mulher", 3):
        print(" ", processo)

    # A prova de que o retorno compõe: junta-se com extend, e cada processo
    # continua sabendo de que assunto veio.
    print()
    print("juntando os dois assuntos numa lista só:")
    todos = []
    for nome in listar_assuntos("violência doméstica"):
        todos.extend(buscar_processos(nome, 2))
    print(" ", len(todos), "processos")
    for processo in todos:
        print(f"  {processo['numero_processo']}  {processo['assunto']}")

    # O processo mais movimentado que encontramos no índice: 613 andamentos.
    print()
    print("ultimos_movimentos('0018683-81.2020.8.19.0066', 5):")
    for m in ultimos_movimentos("0018683-81.2020.8.19.0066", 5):
        print(f"  {m['data']}  {m['grau']}  [{m['codigo']}] {m['movimento']}")