#!/usr/bin/env python3
"""
RAG Builder usando LangChain + HuggingFaceEmbeddings (CPU)
Versão CORRIGIDA com logging detalhado para debug
"""

import os
import sys
import json
import argparse
from pathlib import Path
from typing import List, Dict
import logging
import traceback

from langchain.docstore.document import Document
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def load_jsonl(filepath: str) -> List[Dict]:
    """Carrega registros de arquivo JSONL"""
    records = []
    try:
        logger.info(f"📂 Abrindo arquivo: {filepath}")

        if not os.path.exists(filepath):
            raise FileNotFoundError(f"Arquivo não encontrado: {filepath}")

        file_size = os.path.getsize(filepath)
        logger.info(f"📊 Tamanho do arquivo: {file_size:,} bytes")

        with open(filepath, 'r', encoding='utf-8') as f:
            for i, line in enumerate(f, 1):
                if line.strip():
                    try:
                        records.append(json.loads(line))
                    except json.JSONDecodeError as e:
                        logger.warning(f"⚠️ Linha {i} inválida: {e}")
                        continue

                if i % 50000 == 0:
                    logger.info(f"  Lidas {i:,} linhas...")

        logger.info(f"✅ {len(records):,} registros carregados")
        return records

    except Exception as e:
        logger.error(f"❌ Erro ao carregar JSONL: {type(e).__name__}: {e}")
        raise

def create_documents(records: List[Dict]) -> List[Document]:
    """Converte registros JSONL em Documents do LangChain"""
    documents = []
    skipped = 0

    logger.info("📄 Criando Documents do LangChain...")

    for i, record in enumerate(records, 1):
        doc_id = record.get('id', f'unknown_{i}')
        ementa = record.get('ementa', '')

        if not ementa or not ementa.strip():
            skipped += 1
            continue

        doc = Document(
            page_content=ementa,
            metadata={
                'id': str(doc_id),
                'source': 'tjpr_jurisprudencia'
            }
        )
        documents.append(doc)

        if i % 50000 == 0:
            logger.info(f"  Processados {i:,}/{len(records):,} registros...")

    logger.info(f"✅ {len(documents):,} documentos criados")
    if skipped > 0:
        logger.info(f"⚠️ {skipped:,} registros sem ementa (ignorados)")

    return documents

def build_vectorstore(
    input_file: str,
    output_dir: str = '/app/faiss_index',
    model_name: str = 'sentence-transformers/all-MiniLM-L6-v2',
    batch_size: int = 16
):
    """
    Constrói FAISS vector store com logging detalhado
    """
    try:
        logger.info("="*80)
        logger.info("🚀 RAG Builder - LangChain + FAISS (CPU)")
        logger.info("="*80)
        logger.info(f"📅 Início: {__import__('datetime').datetime.now()}")
        logger.info("")

        # ====================================================================
        # PASSO 1: Carregar JSONL
        # ====================================================================
        logger.info("PASSO 1/5: Carregando arquivo JSONL")
        logger.info("-"*80)
        records = load_jsonl(input_file)

        if not records:
            raise ValueError("❌ Nenhum registro encontrado no arquivo JSONL!")

        logger.info("")

        # ====================================================================
        # PASSO 2: Criar Documents
        # ====================================================================
        logger.info("PASSO 2/5: Convertendo para Documents do LangChain")
        logger.info("-"*80)
        documents = create_documents(records)

        if not documents:
            raise ValueError("❌ Nenhum documento válido criado!")

        logger.info("")

        # ====================================================================
        # PASSO 3: Inicializar Embeddings
        # ====================================================================
        logger.info("PASSO 3/5: Inicializando HuggingFaceEmbeddings")
        logger.info("-"*80)
        logger.info(f"  🤖 Modelo: {model_name}")
        logger.info(f"  💻 Device: CPU")
        logger.info(f"  📦 Batch size: {batch_size}")
        logger.info(f"  📏 Dimensões: 384")
        logger.info("")
        logger.info("  Baixando modelo do HuggingFace Hub...")
        logger.info("  (Primeira vez pode demorar ~1min)")

        embeddings = HuggingFaceEmbeddings(
            model_name=model_name,
            model_kwargs={'device': 'cpu'},
            encode_kwargs={
                'batch_size': batch_size,
                'show_progress_bar': True,
                'normalize_embeddings': True
            }
        )

        logger.info("✅ Embeddings inicializados com sucesso")
        logger.info("")

        # ====================================================================
        # PASSO 4: Construir FAISS Index
        # ====================================================================
        logger.info("PASSO 4/5: Construindo FAISS vector store")
        logger.info("-"*80)
        logger.info(f"  📊 Total de documentos: {len(documents):,}")
        logger.info(f"  ⏱️  Tempo estimado: ~{len(documents) // 500:.0f} minutos")
        logger.info("")
        logger.info("  Processando embeddings (isso vai demorar)...")

        import time
        start_time = time.time()

        vectorstore = FAISS.from_documents(documents, embeddings)

        elapsed = time.time() - start_time
        logger.info(f"✅ FAISS index construído em {elapsed:.1f} segundos")
        logger.info(f"  📈 Velocidade: {len(documents) / elapsed:.0f} docs/seg")
        logger.info("")

        # ====================================================================
        # PASSO 5: Salvar Index
        # ====================================================================
        logger.info("PASSO 5/5: Salvando FAISS index")
        logger.info("-"*80)
        logger.info(f"  💾 Diretório: {output_dir}")

        os.makedirs(output_dir, exist_ok=True)
        vectorstore.save_local(output_dir)

        # Verificar tamanho do index salvo
        index_files = list(Path(output_dir).glob('*'))
        total_size = sum(f.stat().st_size for f in index_files if f.is_file())

        logger.info(f"✅ Index salvo com sucesso")
        logger.info(f"  📁 Arquivos: {len(index_files)}")
        logger.info(f"  📊 Tamanho total: {total_size / 1024 / 1024:.1f} MB")
        logger.info("")

        # ====================================================================
        # RESUMO FINAL
        # ====================================================================
        logger.info("="*80)
        logger.info("✅ BUILD COMPLETO!")
        logger.info("="*80)
        logger.info(f"📊 Documentos indexados: {len(documents):,}")
        logger.info(f"📁 Index salvo em: {output_dir}")
        logger.info(f"🤖 Modelo: {model_name}")
        logger.info(f"📏 Dimensões: 384")
        logger.info(f"⏱️  Tempo total: {time.time() - start_time:.1f}s")
        logger.info("="*80)

        return vectorstore

    except ImportError as e:
        logger.error("")
        logger.error("="*80)
        logger.error("❌ ERRO DE IMPORTAÇÃO")
        logger.error("="*80)
        logger.error(f"Módulo faltando: {e}")
        logger.error("")
        logger.error("SOLUÇÃO:")
        logger.error("  Verificar se requirements.txt inclui:")
        logger.error("  - langchain==0.1.11")
        logger.error("  - langchain-community==0.0.24")
        logger.error("  - sentence-transformers==2.5.1")
        logger.error("  - torch==2.2.0")
        logger.error("  - transformers==4.37.2")
        logger.error("  - faiss-cpu==1.8.0")
        logger.error("="*80)
        raise

    except FileNotFoundError as e:
        logger.error("")
        logger.error("="*80)
        logger.error("❌ ARQUIVO NÃO ENCONTRADO")
        logger.error("="*80)
        logger.error(f"{e}")
        logger.error("")
        logger.error("SOLUÇÃO:")
        logger.error("  Verificar se setup.py criou /tmp/filtered.jsonl")
        logger.error("  Verificar se filter_fields.py foi executado")
        logger.error("="*80)
        raise

    except Exception as e:
        logger.error("")
        logger.error("="*80)
        logger.error("❌ ERRO DURANTE BUILD FAISS")
        logger.error("="*80)
        logger.error(f"Tipo: {type(e).__name__}")
        logger.error(f"Mensagem: {str(e)}")
        logger.error("")
        logger.error("Traceback completo:")
        logger.error("-"*80)
        logger.error(traceback.format_exc())
        logger.error("="*80)
        raise

def main():
    parser = argparse.ArgumentParser(
        description='Build FAISS vector store from JSONL'
    )
    parser.add_argument('--input', required=True, help='Input JSONL file')
    parser.add_argument('--output', default='/app/faiss_index', help='Output directory')
    parser.add_argument(
        '--model',
        default='sentence-transformers/all-MiniLM-L6-v2',
        help='HuggingFace model name'
    )
    parser.add_argument('--batch-size', type=int, default=16, help='Batch size')

    args = parser.parse_args()

    build_vectorstore(
        input_file=args.input,
        output_dir=args.output,
        model_name=args.model,
        batch_size=args.batch_size
    )

if __name__ == '__main__':
    main()