# --- 1. IMPORTS ---
# Imports do seu train.py e do Gradio
import os
import glob
import json
import csv
import numpy as np
from sentence_transformers import SentenceTransformer
import zipfile
import xml.etree.ElementTree as ET
import gradio as gr
import shutil

# --- 2. CONFIGURAÇÕES E CONSTANTES ---
# Caminhos relativos são melhores para portabilidade no Hugging Face Spaces
DATA_DIR = "dados"
EXTRACT_DIR = os.path.join(DATA_DIR, "dados_extraidos")

# ATIVA O ARMAZENAMENTO PERSISTENTE NAS CONFIGURAÇÕES DO SEU SPACE!
# Se ativado, mude o caminho para algo como "/data/meus_embeddings.npy"
OUTPUT_FILENAME = "meus_embeddings_e5_large.npy" 

# --- 3. SUAS FUNÇÕES DE PROCESSAMENTO (DO TRAIN.PY) ---
# Copiamos suas funções de ajuda diretamente para cá.

def setup_data():
    """Descompacta os arquivos .zip e retorna o diretório de processamento."""
    os.makedirs(EXTRACT_DIR, exist_ok=True)
    zip_files = glob.glob(os.path.join(DATA_DIR, "*.zip")) # Simplificado para buscar zips na pasta 'dados'
    if not zip_files:
        print("Nenhum arquivo .zip encontrado, usando o diretório de dados principal.")
        return DATA_DIR
    for zip_path in zip_files:
        with zipfile.ZipFile(zip_path, 'r') as zf:
            zf.extractall(EXTRACT_DIR)
    return EXTRACT_DIR

def xml_to_dict(element):
    """Converte um elemento XML para um dicionário Python."""
    d = {}
    for child in element:
        child_dict = xml_to_dict(child)
        if child.tag in d:
            if not isinstance(d[child.tag], list):
                d[child.tag] = [d[child.tag]]
            d[child.tag].append(child_dict)
        else:
            d[child.tag] = child_dict
    if not d:
        return element.text
    return d

def serialize_item_to_text(item_dict):
    """Converte um dicionário (de JSON, CSV, etc.) para uma string de texto plano."""
    parts = []
    if not isinstance(item_dict, dict):
        return str(item_dict)
    for key, value in item_dict.items():
        if isinstance(value, dict):
            nested_text = serialize_item_to_text(value)
            parts.append(f"{key} ({nested_text})")
        elif isinstance(value, list):
            list_str = ', '.join([serialize_item_to_text(i) for i in value])
            parts.append(f"{key}: [{list_str}]")
        else:
            parts.append(f"{key}: {value}")
    return ", ".join(parts)

# --- 4. FUNÇÃO ORQUESTRADORA (LÓGICA PRINCIPAL) ---
# Esta função substitui a sua função `main()` e é chamada pelo Gradio.
# Ela usa `yield` para enviar atualizações de progresso para a interface.

def run_full_process():
    """Executa o pipeline completo e envia o progresso para a UI."""
    
    # --- ETAPA 1: SETUP E PROCESSAMENTO DE ARQUIVOS ---
    yield "Iniciando... Descompactando arquivos..."
    process_dir = setup_data()
    
    csv.field_size_limit(10_000_000)
    all_files = glob.glob(os.path.join(process_dir, "**/*.json"), recursive=True) + \
                glob.glob(os.path.join(process_dir, "**/*.csv"), recursive=True) + \
                glob.glob(os.path.join(process_dir, "**/*.xml"), recursive=True)
    
    yield f"🔎 Encontrados {len(all_files)} arquivos para processar."

    documents = []
    for idx, filepath in enumerate(all_files):
        try:
            # Mostra o progresso na interface em vez de usar tqdm
            yield f"Processando arquivo {idx + 1}/{len(all_files)}: {os.path.basename(filepath)}"
            
            if filepath.endswith('.json'):
                with open(filepath, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    if isinstance(data, list):
                        for item in data: documents.append(serialize_item_to_text(item))
                    else:
                        documents.append(serialize_item_to_text(data))
            elif filepath.endswith('.csv'):
                with open(filepath, 'r', encoding='utf-8') as f:
                    reader = csv.DictReader(f)
                    for row in reader: documents.append(serialize_item_to_text(row))
            elif filepath.endswith('.xml'):
                tree = ET.parse(filepath)
                root = tree.getroot()
                xml_dict = {root.tag: xml_to_dict(root)}
                documents.append(serialize_item_to_text(xml_dict))
        except Exception as e:
            yield f"⚠️ Erro ao processar {os.path.basename(filepath)}: {e}"

    yield f"Processamento de arquivos concluído! {len(documents)} documentos criados."
    if not documents:
        yield "Nenhum documento encontrado para gerar embeddings. Processo encerrado."
        return

    # --- ETAPA 2: GERAÇÃO DE EMBEDDINGS ---
    yield "Carregando modelo de alta performance: intfloat/multilingual-e5-large..."
    
    # Use um cache dentro do seu Space para não baixar o modelo toda vez
    cache_path = './model_cache'
    os.makedirs(cache_path, exist_ok=True)
    model = SentenceTransformer('intfloat/multilingual-