scripts/setup_neon.py

"""
Script interativo para configurar Neon para RAG Template.

Facilita a configuracao do DATABASE_URL e testa a conexao.
"""

import os
import sys
from urllib.parse import quote_plus


def print_header():
    """Imprime o cabecalho do script."""
    print("=" * 60)
    print(" Setup Neon para RAG Template")
    print("=" * 60)
    print()


def get_connection_string():
    """Solicita a connection string do Neon."""
    print("1. Obter Connection String")
    print("   - Acesse seu projeto no Neon")
    print("   - Va para Connection Details")
    print("   - Copie a Connection String completa")
    print("   - Formato: postgresql://user:pass@ep-xxx.neon.tech/neondb?sslmode=require")
    print()
    print("   DICA: Voce pode colar a string completa aqui")
    print()

    conn_string = input("Connection String: ").strip()
    return conn_string


def validate_connection_string(conn_string: str) -> bool:
    """Valida se a connection string tem formato correto."""
    if not conn_string.startswith("postgresql://"):
        return False
    if "neon.tech" not in conn_string:
        return False
    return True


def ensure_ssl_mode(conn_string: str) -> str:
    """Garante que sslmode=require esta presente."""
    if "sslmode=require" not in conn_string:
        separator = "?" if "?" not in conn_string else "&"
        conn_string += f"{separator}sslmode=require"
        print("   [INFO] Adicionado sslmode=require a URL")
    return conn_string


def save_to_env(database_url: str) -> bool:
    """Salva DATABASE_URL no arquivo .env."""
    print("\n2. Salvar no .env")

    # Verificar se .env ja existe
    env_path = ".env"
    if os.path.exists(env_path):
        print(f"   Arquivo .env encontrado")
        overwrite = input("   Sobrescrever DATABASE_URL existente? (s/n): ").strip().lower()
        if overwrite != 's':
            print("   Pulando salvamento no .env")
            return False

    # Ler conteudo existente
    existing_content = ""
    if os.path.exists(env_path):
        with open(env_path, 'r') as f:
            lines = f.readlines()
            # Remover linha DATABASE_URL existente
            existing_content = "".join(
                line for line in lines if not line.startswith("DATABASE_URL=")
            )

    # Adicionar nova DATABASE_URL
    with open(env_path, 'w') as f:
        f.write(existing_content)
        if existing_content and not existing_content.endswith('\n'):
            f.write('\n')
        f.write(f"DATABASE_URL={database_url}\n")

    print(f"   DATABASE_URL salvo em {env_path}")
    return True


def test_connection(database_url: str) -> bool:
    """Testa a conexao com o banco."""
    print("\n3. Testar Conexao")
    test = input("   Testar conexao agora? (s/n): ").strip().lower()

    if test != 's':
        print("   Pulando teste de conexao")
        return False

    try:
        # Importar aqui para nao quebrar se dependencias nao instaladas
        import psycopg
        from psycopg import sql

        print("   Conectando ao Neon...")

        with psycopg.connect(database_url) as conn:
            with conn.cursor() as cur:
                # Testar conexao
                cur.execute("SELECT version();")
                version = cur.fetchone()[0]
                print(f"   Conexao bem-sucedida!")
                print(f"   PostgreSQL version: {version[:50]}...")

                # Verificar extensao pgvector
                cur.execute(
                    "SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'vector');"
                )
                has_vector = cur.fetchone()[0]

                if has_vector:
                    print("   Extensao pgvector: INSTALADA")
                else:
                    print("   Extensao pgvector: NAO ENCONTRADA")
                    print("\n   Para instalar, execute no SQL Editor do Neon:")
                    print("   CREATE EXTENSION vector;")

                # Verificar storage usado
                cur.execute(
                    "SELECT pg_size_pretty(pg_database_size(current_database()));"
                )
                db_size = cur.fetchone()[0]
                print(f"   Database size: {db_size}")

        return True

    except ImportError:
        print("   ERRO: psycopg nao instalado")
        print("   Execute: pip install psycopg[binary]")
        return False

    except Exception as e:
        print(f"   ERRO ao conectar: {e}")
        print("\n   Verifique:")
        print("   - Connection string correta")
        print("   - Projeto Neon ativo")
        print("   - Firewall nao bloqueando conexao")
        return False


def print_tips():
    """Imprime dicas sobre o Neon."""
    print("\n" + "=" * 60)
    print(" Dicas Neon")
    print("=" * 60)
    print("\n  Free Tier:")
    print("  - 10GB storage")
    print("  - 100 compute hours/mes")
    print("  - 10 projetos")
    print("  - 10 branches por projeto")
    print("\n  Branching:")
    print("  - Crie branches para testes sem afetar producao")
    print("  - Cada branch tem sua propria connection string")
    print("\n  Performance:")
    print("  - Use connection pooling (?pooler=true na URL)")
    print("  - Neon pausa automaticamente apos 5min inatividade")
    print("  - Cold start geralmente <2s")


def print_next_steps(success: bool):
    """Imprime proximos passos."""
    print("\n" + "=" * 60)
    print(" Proximos Passos")
    print("=" * 60)

    if success:
        print("\n  Conexao configurada com sucesso!")
        print("\n  1. Habilite pgvector (se nao habilitado):")
        print("     - Va ao SQL Editor no Neon")
        print("     - Execute: CREATE EXTENSION vector;")
        print("\n  2. Configure HF_TOKEN no .env:")
        print("     HF_TOKEN=seu_token_huggingface")
        print("\n  3. Execute o app:")
        print("     python app.py")
        print("\n  4. Acesse: http://localhost:7860")
    else:
        print("\n  Configure DATABASE_URL manualmente no .env:")
        print("  DATABASE_URL=postgresql://user:pass@ep-xxx.neon.tech/neondb?sslmode=require")
        print("\n  Ou execute este script novamente")

    print("\n  Documentacao completa: docs/NEON_SETUP.md")
    print("=" * 60)


def main():
    """Funcao principal."""
    print_header()

    # Coletar connection string
    conn_string = get_connection_string()

    if not conn_string:
        print("ERRO: Connection string vazia")
        sys.exit(1)

    # Validar formato
    if not validate_connection_string(conn_string):
        print("ERRO: Connection string invalida")
        print("Formato esperado: postgresql://user:pass@ep-xxx.neon.tech/neondb")
        sys.exit(1)

    # Garantir SSL mode
    database_url = ensure_ssl_mode(conn_string)

    print("\n" + "=" * 60)
    print(" DATABASE_URL Configurado")
    print("=" * 60)
    print(f"\n{database_url}\n")

    # Salvar no .env
    saved = save_to_env(database_url)

    # Testar conexao
    success = test_connection(database_url)

    # Dicas
    if success:
        print_tips()

    # Proximos passos
    print_next_steps(success and saved)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\nSetup cancelado pelo usuario")
        sys.exit(0)