#!/bin/bash

# Script para processar documentos RAG e gerar embeddings
# Museu Virtual Mario

set -e

# Configurações
RAG_DIR="/opt/museu-data/rag"
DOCUMENTS_DIR="$RAG_DIR/documents"
EMBEDDINGS_DIR="$RAG_DIR/embeddings"
VERSIONS_DIR="$RAG_DIR/versions"
LOG_FILE="/var/log/museu/rag_processing.log"
PYTHON_ENV="/opt/rag-env"

# Cores
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

echo -e "${BLUE}🧠 Processamento RAG - Museu Virtual Mario${NC}"
echo "=================================================="

# Criar diretórios de log
sudo mkdir -p "$(dirname "$LOG_FILE")"
sudo touch "$LOG_FILE"

# Função de log
log() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') - $1" | sudo tee -a "$LOG_FILE"
}

# Verificar ambiente Python
echo ""
echo -e "${BLUE}🐍 Verificando ambiente Python...${NC}"
if [ ! -d "$PYTHON_ENV" ]; then
    echo -e "${YELLOW}⚠️ Criando ambiente virtual Python...${NC}"
    sudo python3 -m venv "$PYTHON_ENV"
    sudo "$PYTHON_ENV/bin/pip" install --upgrade pip
    sudo "$PYTHON_ENV/bin/pip" install sentence-transformers torch numpy pandas PyYAML markdown
fi

# Ativar ambiente virtual
source "$PYTHON_ENV/bin/activate"

# Verificar dependências
echo ""
echo -e "${BLUE}📦 Verificando dependências...${NC}"
python3 -c "
import sentence_transformers
import torch
import numpy
import pandas
import yaml
import markdown
print('✅ Todas as dependências estão instaladas')
" || {
    echo -e "${RED}❌ Falha na verificação de dependências${NC}"
    log "ERRO: Dependências Python não encontradas"
    exit 1
}

# Função para processar documento
process_document() {
    local file="$1"
    local filename=$(basename "$file" .md)
    
    echo "📄 Processando: $filename"
    
    # Extrair metadados e conteúdo
    python3 << EOF
import yaml
import markdown
import json
import re
from pathlib import Path

# Ler arquivo
with open('$file', 'r', encoding='utf-8') as f:
    content = f.read()

# Separar metadados
metadata_match = re.match(r'^---\n(.*?)\n---\n(.*)$', content, re.DOTALL)
if metadata_match:
    metadata = yaml.safe_load(metadata_match.group(1))
    markdown_content = metadata_match.group(2)
else:
    print(f"ERRO: Metadados não encontrados em {file}")
    exit(1)

# Converter markdown para texto puro
html = markdown.markdown(markdown_content)
text_content = re.sub(r'<[^>]+>', '', html)

# Estruturar dados
document_data = {
    'filename': '$filename',
    'metadata': metadata,
    'content': text_content,
    'processed_at': '$(date +%Y-%m-%d_%H:%M:%S)',
    'word_count': len(text_content.split()),
    'char_count': len(text_content)
}

# Salvar JSON processado
output_dir = Path('$EMBEDDINGS_DIR/processed')
output_dir.mkdir(parents=True, exist_ok=True)

with open(output_dir / f'{filename}.json', 'w', encoding='utf-8') as f:
    json.dump(document_data, f, ensure_ascii=False, indent=2)

print(f"✅ Processado: {filename} ({document_data['word_count']} palavras)")
EOF
}

# Função para gerar embeddings
generate_embeddings() {
    local processed_file="$1"
    local filename=$(basename "$processed_file" .json)
    
    echo "🧠 Gerando embeddings: $filename"
    
    python3 << EOF
import json
import numpy as np
from sentence_transformers import SentenceTransformer
from pathlib import Path

# Carregar modelo
model = SentenceTransformer('all-MiniLM-L6-v2')

# Ler documento processado
with open('$processed_file', 'r', encoding='utf-8') as f:
    doc_data = json.load(f)

# Gerar embeddings
content = doc_data['content']
embeddings = model.encode(content)

# Salvar embeddings
embeddings_dir = Path('$EMBEDDINGS_DIR/vectors')
embeddings_dir.mkdir(parents=True, exist_ok=True)

embedding_data = {
    'filename': doc_data['filename'],
    'metadata': doc_data['metadata'],
    'embeddings': embeddings.tolist(),
    'embedding_dim': len(embeddings),
    'processed_at': doc_data['processed_at']
}

with open(embeddings_dir / f'{filename}.npy', 'wb') as f:
    np.save(f, embeddings)

with open(embeddings_dir / f'{filename}_meta.json', 'w', encoding='utf-8') as f:
    json.dump(embedding_data, f, ensure_ascii=False, indent=2)

print(f"✅ Embeddings gerados: {filename} (dimensão: {len(embeddings)})")
EOF
}

# Função para criar índice
create_index() {
    echo "📚 Criando índice de busca..."
    
    python3 << EOF
import json
import numpy as np
from pathlib import Path
from sklearn.metrics.pairwise import cosine_similarity

# Carregar todos os embeddings
embeddings_dir = Path('$EMBEDDINGS_DIR/vectors')
embeddings_data = []

for meta_file in embeddings_dir.glob('*_meta.json'):
    with open(meta_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    vector_file = meta_file.with_suffix('.npy')
    embeddings = np.load(vector_file)
    
    embeddings_data.append({
        'filename': data['filename'],
        'metadata': data['metadata'],
        'embeddings': embeddings,
        'embedding_dim': data['embedding_dim']
    })

# Criar índice
index_data = {
    'documents': embeddings_data,
    'total_documents': len(embeddings_data),
    'embedding_dim': embeddings_data[0]['embedding_dim'] if embeddings_data else 0,
    'created_at': '$(date +%Y-%m-%d_%H:%M:%S)',
    'categories': list(set(doc['metadata'].get('category', 'unknown') for doc in embeddings_data))
}

# Salvar índice
index_dir = Path('$EMBEDDINGS_DIR/index')
index_dir.mkdir(parents=True, exist_ok=True)

with open(index_dir / 'rag_index.json', 'w', encoding='utf-8') as f:
    json.dump(index_data, f, ensure_ascii=False, indent=2)

print(f"✅ Índice criado: {len(embeddings_data)} documentos")
EOF
}

# Processar todos os documentos
echo ""
echo -e "${BLUE}📄 Processando documentos...${NC}"
processed_count=0

for category_dir in "$DOCUMENTS_DIR"/*; do
    if [ -d "$category_dir" ]; then
        category=$(basename "$category_dir")
        echo "📂 Processando categoria: $category"
        
        for doc_file in "$category_dir"/*.md; do
            if [ -f "$doc_file" ]; then
                process_document "$doc_file"
                ((processed_count++))
            fi
        done
    fi
done

echo -e "${GREEN}✅ $processed_count documentos processados${NC}"

# Gerar embeddings
echo ""
echo -e "${BLUE}🧠 Gerando embeddings...${NC}"
embeddings_count=0

for processed_file in "$EMBEDDINGS_DIR/processed"/*.json; do
    if [ -f "$processed_file" ]; then
        generate_embeddings "$processed_file"
        ((embeddings_count++))
    fi
done

echo -e "${GREEN}✅ $embeddings_count embeddings gerados${NC}"

# Criar índice final
echo ""
create_index

# Gerar relatório
echo ""
echo -e "${BLUE}📊 Gerando relatório...${NC}"
report_file="$EMBEDDINGS_DIR/report_$(date +%Y%m%d_%H%M%S).json"

python3 << EOF
import json
from pathlib import Path

# Coletar estatísticas
processed_dir = Path('$EMBEDDINGS_DIR/processed')
vectors_dir = Path('$EMBEDDINGS_DIR/vectors')

processed_files = list(processed_dir.glob('*.json'))
vector_files = list(vectors_dir.glob('*.npy'))

# Calcular estatísticas
total_words = 0
total_chars = 0
categories = {}

for file in processed_files:
    with open(file, 'r', encoding='utf-8') as f:
        data = json.load(f)
        total_words += data['word_count']
        total_chars += data['char_count']
        
        category = data['metadata'].get('category', 'unknown')
        categories[category] = categories.get(category, 0) + 1

report = {
    'processing_date': '$(date +%Y-%m-%d_%H:%M:%S)',
    'statistics': {
        'total_documents': len(processed_files),
        'total_embeddings': len(vector_files),
        'total_words': total_words,
        'total_characters': total_chars,
        'categories': categories,
        'average_words_per_doc': total_words / len(processed_files) if processed_files else 0
    },
    'directories': {
        'documents': str(processed_dir),
        'embeddings': str(vectors_dir),
        'index': str(Path('$EMBEDDINGS_DIR/index'))
    }
}

with open('$report_file', 'w', encoding='utf-8') as f:
    json.dump(report, f, ensure_ascii=False, indent=2)

print(f"✅ Relatório gerado: {report_file.name}")
EOF

echo -e "${GREEN}✅ Processamento RAG concluído!${NC}"
echo ""
echo "📊 Estatísticas:"
echo "  - Documentos processados: $processed_count"
echo "  - Embeddings gerados: $embeddings_count"
echo "  - Índice criado: $EMBEDDINGS_DIR/index/rag_index.json"
echo "  - Relatório: $report_file"
echo ""
echo "📝 Log completo: $LOG_FILE"
echo ""
echo "🔄 Para atualização incremental:"
echo "  ./scripts/process-rag.sh --update"

log "PROCESSAMENTO CONCLUÍDO: $processed_count documentos, $embeddings_count embeddings"

# Desativar ambiente virtual
deactivate
