#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Helper: Criar chunks de dados para beeRoot
"""

import json
import tarfile
import os
from pathlib import Path

def create_chunk(documents, chunk_id, output_dir="chunks_dados"):
    """
    Cria um chunk .tar.gz a partir de documentos

    Args:
        documents: Lista de dicts com 'Id' e 'ementa'
        chunk_id: Número do chunk (1, 2, 3...)
        output_dir: Diretório de saída
    """
    # Limita a 500 docs por chunk
    documents = documents[:500]

    # Nomes
    chunk_name = f"chunk_dados_{chunk_id:06d}"
    jsonl_path = f"{chunk_name}.jsonl"
    tar_path = f"{output_dir}/{chunk_name}.tar.gz"

    # Cria output dir
    Path(output_dir).mkdir(exist_ok=True)

    # Cria JSONL
    with open(jsonl_path, 'w', encoding='utf-8') as f:
        for doc in documents:
            # Valida campos obrigatórios
            if 'Id' not in doc or 'ementa' not in doc:
                print(f"⚠️  Documento sem Id ou ementa: {doc}")
                continue

            f.write(json.dumps(doc, ensure_ascii=False) + '\n')

    # Compacta
    with tarfile.open(tar_path, 'w:gz') as tar:
        tar.add(jsonl_path, arcname='jurisprudencias.jsonl')

    # Remove temp
    os.remove(jsonl_path)

    print(f"✅ Criado: {tar_path} ({len(documents)} docs)")
    return tar_path

# Exemplo de uso
if __name__ == '__main__':
    # Seus documentos
    docs = [
        {"Id": "001", "ementa": "Decisão sobre FGTS..."},
        {"Id": "002", "ementa": "Recurso sobre aposentadoria..."},
        # ... até 500 documentos
    ]

    # Cria chunk
    create_chunk(docs, chunk_id=1, output_dir="meus_chunks")

    # Para múltiplos chunks
    all_docs = [...] # Sua lista completa

    chunk_size = 500
    for i in range(0, len(all_docs), chunk_size):
        chunk_docs = all_docs[i:i+chunk_size]
        chunk_id = (i // chunk_size) + 1
        create_chunk(chunk_docs, chunk_id)