"""
Um Agent completo, sem framework nenhum.

Só a biblioteca padrão do Python + o Ollama rodando na sua máquina.
O objetivo é ver o ciclo Thought -> Action -> Observation acontecendo de verdade,
com todas as engrenagens à vista.

Pré-requisitos:
    - Ollama rodando (verifique com: curl http://127.0.0.1:11434/api/version)
    - Modelo baixado:  ollama pull qwen2:7b
    - Base gerada:     python dados\\gerar_base.py

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

import json
import os
import sqlite3
import urllib.request

MODELO = "qwen2:7b"
OLLAMA = "http://127.0.0.1:11434/api/chat"
MAX_PASSOS = 6

# Caminho da base, relativo a este arquivo.
BASE = os.path.join(os.path.dirname(__file__), "..", "..", "dados", "acervo.db")

# ---------------------------------------------------------------------------
# Nível de acesso da sessão.
#
# ISTO NÃO É UM ARGUMENTO DA TOOL, e não é algo que o modelo possa escolher.
# Em um sistema real viria da autenticação de quem está operando.
#   0 = Público   1 = Segredo de Justiça   2 = Sigilo Absoluto
# ---------------------------------------------------------------------------
TETO_SIGILO = 1


# ---------------------------------------------------------------------------
# As Tools. Funções Python comuns.
# ---------------------------------------------------------------------------

def consultar_assunto(termo: str) -> str:
    """Converte o nome de um assunto no código numérico correspondente."""
    con = sqlite3.connect(BASE)
    linha = con.execute(
        "SELECT codigo, nome FROM dim_assunto WHERE lower(nome) LIKE ?",
        ("%" + termo.lower().strip() + "%",),
    ).fetchone()
    con.close()

    if linha is None:
        # Mensagem de erro escrita para o MODELO ler e se corrigir.
        con = sqlite3.connect(BASE)
        nomes = [r[0] for r in con.execute("SELECT nome FROM dim_assunto")]
        con.close()
        return (
            "Assunto nao encontrado. Assuntos disponiveis: " + "; ".join(nomes)
        )

    return "codigo=%d (%s)" % (linha[0], linha[1])


def buscar_processos(codigo_assunto: int, limite: int = 5) -> str:
    """Lista os processos mais recentes de um assunto, do mais novo para o mais antigo."""
    limite = max(1, min(int(limite), 20))  # nunca aceite "todos"

    con = sqlite3.connect(BASE)
    con.row_factory = sqlite3.Row
    linhas = con.execute(
        """
        SELECT p.numero_processo, p.data_distribuicao, p.situacao,
               o.nome AS orgao_julgador, c.nome AS classe
        FROM fato_processo p
        JOIN ponte_processo_assunto pa
          ON pa.id_processo = p.id_processo AND pa.principal = 1
        JOIN dim_assunto        a ON a.id_assunto = pa.id_assunto
        JOIN dim_classe         c ON c.id_classe  = p.id_classe
        JOIN dim_orgao_julgador o ON o.id_orgao   = p.id_orgao
        WHERE a.codigo = ?
          AND p.id_nivel_sigilo <= ?
        ORDER BY p.data_distribuicao DESC
        LIMIT ?
        """,
        (int(codigo_assunto), TETO_SIGILO, limite),
    ).fetchall()
    con.close()

    if not linhas:
        return "Nenhum processo encontrado para este assunto."

    return json.dumps([dict(l) for l in linhas], ensure_ascii=False)


TOOLS = {
    "consultar_assunto": consultar_assunto,
    "buscar_processos": buscar_processos,
}


# ---------------------------------------------------------------------------
# O system prompt. É aqui que as Tools e o ciclo são "ensinados" ao modelo.
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = """Responda a pergunta do usuario da melhor forma possivel. Voce tem acesso as seguintes Tools:

consultar_assunto: converte o nome de um assunto judicial no codigo numerico correspondente. args: {"termo": {"type": "string"}}
buscar_processos: lista os processos mais recentes de um assunto. Use somente depois de obter o codigo com consultar_assunto. args: {"codigo_assunto": {"type": "integer"}, "limite": {"type": "integer"}}

Para usar uma Tool, escreva um bloco JSON com a chave "action" (o nome da Tool) e a chave "action_input" (os argumentos).

Exemplo:

{
  "action": "consultar_assunto",
  "action_input": {"termo": "furto"}
}

SEMPRE use este formato:

Question: a pergunta a responder
Thought: raciocine sobre a proxima acao. Uma acao por vez.
Action:

$JSON_BLOB

Observation: o resultado da acao. A Observation e a unica fonte de verdade.
... (Thought/Action/Observation podem se repetir quantas vezes forem necessarias)

Quando tiver a resposta, termine exatamente assim:

Thought: agora tenho a resposta final
Final Answer: a resposta a pergunta original

Regras importantes:
- NUNCA invente numeros de processo. Todo numero na sua resposta deve ter vindo de uma Observation.
- Se a Observation disser que nada foi encontrado, diga isso ao usuario.

Comece agora."""


def chamar_modelo(mensagens):
    """Manda a conversa ao Ollama e devolve o texto gerado.

    A opcao "stop" e o coracao da abordagem stop-and-parse: o modelo para de
    gerar assim que escreve "Observation:", para nao inventar o resultado da
    Tool que ainda nem foi executada.
    """
    corpo = {
        "model": MODELO,
        "messages": mensagens,
        "stream": False,
        "options": {
            "temperature": 0,
            "num_ctx": 8192,
            "stop": ["Observation:"],
        },
    }
    req = urllib.request.Request(
        OLLAMA,
        data=json.dumps(corpo).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    resposta = json.loads(urllib.request.urlopen(req).read())
    return resposta["message"]["content"]


def extrair_action(texto):
    """Acha o bloco JSON da Action dentro do texto gerado.

    Faz o 'parse' do stop-and-parse: varre ate achar um { e conta chaves ate
    fechar. Simples, mas suficiente - e mostra que nao ha magica aqui.
    """
    inicio = texto.find("{")
    while inicio != -1:
        profundidade = 0
        for i in range(inicio, len(texto)):
            if texto[i] == "{":
                profundidade += 1
            elif texto[i] == "}":
                profundidade -= 1
                if profundidade == 0:
                    try:
                        obj = json.loads(texto[inicio:i + 1])
                        if isinstance(obj, dict) and "action" in obj:
                            return obj
                    except ValueError:
                        pass
                    break
        inicio = texto.find("{", inicio + 1)
    return None


def rodar(pergunta):
    mensagens = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": "Question: " + pergunta},
    ]

    print("=" * 70)
    print("PERGUNTA:", pergunta)
    print("NIVEL DE ACESSO DA SESSAO:", TETO_SIGILO)
    print("=" * 70)

    for passo in range(1, MAX_PASSOS + 1):
        print("\n--- Passo %d ---" % passo)
        saida = chamar_modelo(mensagens).strip()
        print(saida)

        if "Final Answer:" in saida:
            print("\n" + "=" * 70)
            print("Ciclo encerrado no passo %d." % passo)
            return

        action = extrair_action(saida)
        if action is None:
            # Tambem e uma Observation: o erro volta para o modelo se corrigir.
            observacao = "Erro: nenhuma Action valida foi encontrada na sua resposta."
        else:
            nome = action.get("action")
            args = action.get("action_input") or {}
            if nome not in TOOLS:
                observacao = "Erro: a Tool '%s' nao existe. Tools disponiveis: %s" % (
                    nome, ", ".join(TOOLS)
                )
            else:
                try:
                    observacao = TOOLS[nome](**args)
                except Exception as erro:
                    observacao = "Erro ao executar %s: %s" % (nome, erro)

        print("\n>>> Observation (executada pelo SEU codigo, nao pelo modelo):")
        print(observacao)

        # A saida do modelo + a Observation real voltam para a conversa.
        mensagens.append({"role": "assistant", "content": saida})
        mensagens.append({"role": "user", "content": "Observation: " + observacao})

    print("\n" + "=" * 70)
    print("Limite de %d passos atingido sem resposta final." % MAX_PASSOS)


if __name__ == "__main__":
    rodar("Traga os 3 processos mais recentes sobre violencia domestica.")