#!/usr/bin/env python3
"""
Setup - Inicialização do sistema
Baixa dados, cria FAISS index, etc
"""

import faiss
import numpy as np
import json
import yaml
from pathlib import Path
from sentence_transformers import SentenceTransformer
import time
import sys

def log(msg):
    """Log com timestamp"""
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] {msg}", flush=True)

def load_model_with_retry(model_name, max_retries=5, delay=30):
    """Carrega modelo com retry automático"""
    for attempt in range(1, max_retries + 1):
        try:
            log(f"⬇️  [{attempt}/{max_retries}] Baixando {model_name}...")
            model = SentenceTransformer(model_name)
            log(f"✅ Modelo carregado!")
            return model
        except Exception as e:
            error_msg = str(e).lower()
            is_network = any(x in error_msg for x in ["resolution", "connection", "timeout", "dns"])

            if is_network and attempt < max_retries:
                log(f"⚠️  Rede falhou. Retry em {delay}s...")
                time.sleep(delay)
            else:
                raise

def main():
    log("="*70)
    log("🐝 beeROOT Para.AI - SETUP")
    log("="*70)

    # Carrega config
    config_path = Path("config.yaml")
    if not config_path.exists():
        log("❌ config.yaml não encontrado!")
        sys.exit(1)

    with open(config_path) as f:
        config = yaml.safe_load(f)

    log(f"📋 Config carregado:")
    log(f"   Modelo: {config.get('model_name', 'paraphrase-multilingual-mpnet-base-v2')}")
    log(f"   Data source: {config.get('data_source', 'local')}")

    # Carrega modelo com retry
    model_name = config.get('model_name', 'paraphrase-multilingual-mpnet-base-v2')
    model = load_model_with_retry(model_name)

    # Prepara diretórios
    faiss_dir = Path("/home/user/app/faiss_index")
    data_dir = Path("/home/user/app/data")
    faiss_dir.mkdir(parents=True, exist_ok=True)
    data_dir.mkdir(parents=True, exist_ok=True)

    # EXEMPLO: Carrega dados de arquivo local
    # Você pode adaptar para baixar do GitHub, etc
    sample_data = [
        "Exemplo de jurisprudência 1: FGTS com correção monetária.",
        "Exemplo de jurisprudência 2: Direitos trabalhistas.",
        "Exemplo de jurisprudência 3: Recurso especial negado."
    ]

    log(f"📄 Dados de exemplo: {len(sample_data)} documentos")

    # Gera embeddings
    log("🔧 Gerando embeddings...")
    embeddings = model.encode(sample_data, show_progress_bar=True)
    embeddings = np.array(embeddings, dtype='float32')

    # Cria FAISS index
    log("🔨 Criando FAISS index...")
    dimension = embeddings.shape[1]
    index = faiss.IndexFlatL2(dimension)
    index.add(embeddings)

    # Salva index
    index_path = faiss_dir / "index.faiss"
    faiss.write_index(index, str(index_path))
    log(f"✅ FAISS index salvo: {index.ntotal} vetores")

    # Salva metadata
    metadata_path = data_dir / "metadata.jsonl"
    with open(metadata_path, 'w') as f:
        for i, text in enumerate(sample_data):
            meta = {
                'text': text,
                'source': 'sample_data',
                'chunk_idx': i
            }
            f.write(json.dumps(meta, ensure_ascii=False) + '\n')

    log(f"✅ Metadata salvo: {len(sample_data)} registros")

    # Marca como pronto
    ready_flag = Path("/tmp/rag_ready")
    ready_flag.touch()

    log("="*70)
    log("✅ SETUP COMPLETO!")
    log("="*70)

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        log(f"❌ ERRO FATAL: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)