import sys
import os
import json
import glob
import asyncio
from langchain_google_genai import ChatGoogleGenerativeAI
from dotenv import load_dotenv
# Garante que o pacote local em src/ seja encontrado antes de qualquer outro
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "src")))
from nexus_rag import NexusPipelineRunner


async def main():
    # Modelo Google Gemini via API — requer GOOGLE_API_KEY no ambiente.
    # gemini-2.5-flash é o mais econômico com suporte a Structured Output;
    # ideal para o tier gratuito (baixo custo por token e alta taxa de req/min).
    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0)

    pipeline = NexusPipelineRunner(
        llm=llm,

        # Tipos de relação aceitos. Qualquer valor fora dessa lista é rejeitado
        # em tempo de parse — a LLM não consegue alucinar uma relação inválida.
        relation_types=["REVOGA", "ALTERA", "ADICIONA", "REGULAMENTA", "COMPLEMENTA"],

        # Dica de formato para o ID das moléculas. Vai direto no prompt da LLM
        # para guiar como ela deve preencher source_id e target_id nos Bonds.
        id_format_hint="NÚMERO/ANO (ex: 361/2026)",

        # Regras de chunking em ordem de prioridade.
        # Cada padrão define um ponto de corte natural no documento.
        # O chunker testa do primeiro ao último — use do mais abrangente ao mais específico.
        chunking_rules=[
            # Nível 0 — blocos estruturais principais (qualquer cabeçalho Markdown h1–h3)
            # Ex: RESOLVE, ANEXO, etc.
            r"^#{1,3}\s+.*",

            # Nível 1 — artigos com tolerância a variações de OCR e formatação
            # Captura: Art. 1, Art. 1º, Artigo 2, ART. 3o, ## Art. 4º, etc.
            r"(?i)^(?:#+\s*)?(?:Art|Artigo)[\.\s]+\d+(?:º|o)?.*",

            # Nível 2 — parágrafos e itens (§ 1º, Parágrafo único, Parágrafo 2)
            r"^(?:-\s*)?(?:§|Parágrafo)[\.\s]*\d+.*",

            # Nível 2b — subitens numerados do tipo "1.1", "2.3.1" (comum em normas técnicas)
            r"^(?:-\s*)?\d+\.\d+.*",

            # Nível 3 — alíneas (a), b), etc.)
            r"^(?:-\s*)?[a-z]\).*"
        ],

        # Padrão de pré-filtro para eleger átomos candidatos à extração.
        # Só os átomos que baterem aqui vão para a LLM — o resto é ignorado,
        # o que economiza bastante em chamadas de API e tempo de GPU.
        #
        # Cobre as formas mais comuns de referência normativa em português:
        #   - Revogações e alterações diretas ("alterado pela", "redação dada pela")
        #   - Inclusões e acréscimos ("incluído pela", "acrescentado pela")
        #   - Vigência e entrada em vigor ("passa a vigorar", "ficam revogados")
        #   - Remissões regulatórias ("estabelecido pela", "regulamentado pela")
        candidate_trigger_pattern=(
            r"(alterad[ao]s?\s+pela|revogad[ao]s?\s+pela|incluíd[ao]s?\s+pela|"
            r"acrescentad[ao]s?(?:\s+pela)?|redação\s+dada\s+pela|fica[m]?\s+revogad[ao]s?|"
            r"passa[m]?\s+a\s+vigorar|altera[r]?\s+(?:a|o|os|dispositivos)|"
            r"incluir\s+(?:o|a|no|na)|estabelecid[ao]\s+pela|regulamentad[ao]\s+pela"
            r").*?(Resolução|Lei|Portaria|Art\.|Decreto)"
        )
    )

    # Diretórios de I/O — gerenciados pelo script, não pela biblioteca
    pdf_dir = "./.documents/pdf"
    md_dir  = "./.documents/markdown"
    out_dir = "./.documents/output"

    os.makedirs(md_dir, exist_ok=True)
    os.makedirs(out_dir, exist_ok=True)

    pdf_files = glob.glob(os.path.join(pdf_dir, "*.pdf"))
    if not pdf_files:
        print(f"Nenhum PDF encontrado em '{pdf_dir}'. Coloque arquivos lá e rode de novo.")
        return

    for pdf_path in pdf_files:
        filename    = os.path.splitext(os.path.basename(pdf_path))[0]
        molecule_id = filename.replace("-", "/")

        print(f"\nProcessando {filename}...")

        molecule, atoms, bonds = await pipeline.process_from_pdf_async(
            pdf_path=pdf_path,
            molecule_id=molecule_id,
            title=f"Documento {molecule_id}",
            aliases=[molecule_id],
        )

        md_dest   = os.path.join(md_dir,  f"{filename}.md")
        json_dest = os.path.join(out_dir, f"{filename}.json")

        # Salva o Markdown gerado pelo Docling para inspeção ou reuso futuro
        with open(md_dest, "w", encoding="utf-8") as f:
            f.write(molecule.content)

        # Serializa as entidades estruturadas — a biblioteca só trabalha em memória,
        # então qualquer persistência fica aqui no script cliente
        with open(json_dest, "w", encoding="utf-8") as f:
            data = {
                "molecule": molecule.model_dump(exclude={"content"}),
                "atoms":    [a.model_dump() for a in atoms],
                "bonds":    [b.model_dump() for b in bonds],
            }
            json.dump(data, f, ensure_ascii=False, indent=2)

        print(f"  Markdown salvo em {md_dest}")
        print(f"  {len(atoms)} átomos e {len(bonds)} ligações salvos em {json_dest}")

    print("\nPipeline concluído.")


if __name__ == "__main__":
    load_dotenv()  # Carrega variáveis de ambiente do arquivo .env, incluindo GOOGLE_API_KEY
    asyncio.run(main())