#!/usr/bin/env python3
"""
Script de migration : Ajout des colonnes SAST dans les tables users ET api_key
Date: 7 novembre 2025
Usage: python scripts/database/migrate_sast_all_tables.py
"""

import os
import sys
from pathlib import Path

# Ajouter le répertoire racine au path
root_dir = Path(__file__).parent.parent.parent
sys.path.insert(0, str(root_dir))

import psycopg2
from dotenv import load_dotenv

# Charger les variables d'environnement
load_dotenv('.env.local')  # Charger le fichier .env.local

# Configuration PostgreSQL
DATABASE_URL = os.getenv('WEB_SENTINEL_POSTGRES_URL') or os.getenv('DATABASE_URL')

if not DATABASE_URL:
    print("❌ Erreur: Variable d'environnement DATABASE_URL non définie")
    print("   Définissez WEB_SENTINEL_POSTGRES_URL ou DATABASE_URL")
    sys.exit(1)

# Mapping des quotas par tier
TIER_QUOTAS = {
    'FREE': {
        'allow_source_scan': False,
        'max_source_files': 0,
        'max_source_size_mb': 0,
        'advanced_rules': False,
    },
    'STARTER': {
        'allow_source_scan': True,
        'max_source_files': 50,
        'max_source_size_mb': 10,
        'advanced_rules': False,
    },
    'PRO': {
        'allow_source_scan': True,
        'max_source_files': 100,
        'max_source_size_mb': 20,
        'advanced_rules': False,
    },
    'ENTERPRISE': {
        'allow_source_scan': True,
        'max_source_files': 500,
        'max_source_size_mb': 100,
        'advanced_rules': True,
    },
    'SYSOP': {
        'allow_source_scan': True,
        'max_source_files': 9999,
        'max_source_size_mb': 500,
        'advanced_rules': True,
    },
}


def migrate_table(cursor, conn, table_name):
    """Migre une table spécifique"""
    print(f"\n🔧 Migration de la table '{table_name}'...")
    
    # Étape 1: Vérifier si les colonnes existent déjà
    cursor.execute(f"""
        SELECT column_name 
        FROM information_schema.columns 
        WHERE table_name = '{table_name}' 
          AND column_name IN ('allow_source_scan', 'max_source_files', 'max_source_size_mb', 'advanced_rules')
    """)
    existing_columns = [row[0] for row in cursor.fetchall()]
    
    if len(existing_columns) == 4:
        print(f"   ℹ️  Colonnes SAST déjà présentes dans '{table_name}'")
        return
    
    # Étape 2: Ajouter les colonnes SAST
    print(f"   📝 Ajout des colonnes SAST dans '{table_name}'...")
    cursor.execute(f"""
        ALTER TABLE {table_name} 
            ADD COLUMN IF NOT EXISTS allow_source_scan BOOLEAN DEFAULT FALSE,
            ADD COLUMN IF NOT EXISTS max_source_files INTEGER DEFAULT 0,
            ADD COLUMN IF NOT EXISTS max_source_size_mb INTEGER DEFAULT 0,
            ADD COLUMN IF NOT EXISTS advanced_rules BOOLEAN DEFAULT FALSE;
    """)
    conn.commit()
    print(f"   ✅ Colonnes ajoutées dans '{table_name}'")
    
    # Étape 3: Récupérer les lignes existantes
    cursor.execute(f"SELECT id, email, tier FROM {table_name}")
    rows = cursor.fetchall()
    print(f"   📌 {len(rows)} lignes trouvées dans '{table_name}'")
    
    if len(rows) == 0:
        return
    
    # Étape 4: Mettre à jour chaque ligne selon son tier
    print(f"   🔄 Mise à jour des quotas...")
    updated_count = 0
    
    for row_id, email, tier in rows:
        tier_upper = tier.upper() if tier else 'FREE'
        quotas = TIER_QUOTAS.get(tier_upper, TIER_QUOTAS['FREE'])
        
        cursor.execute(f"""
            UPDATE {table_name}
            SET 
                allow_source_scan = %s,
                max_source_files = %s,
                max_source_size_mb = %s,
                advanced_rules = %s
            WHERE id = %s
        """, (
            quotas['allow_source_scan'],
            quotas['max_source_files'],
            quotas['max_source_size_mb'],
            quotas['advanced_rules'],
            row_id
        ))
        updated_count += 1
        
        if updated_count <= 5:  # Afficher seulement les 5 premiers
            print(f"   ✓ {email} ({tier_upper}): "
                  f"{quotas['max_source_files']} fichiers, "
                  f"{quotas['max_source_size_mb']} MB")
    
    if updated_count > 5:
        print(f"   ... et {updated_count - 5} autres")
    
    conn.commit()
    print(f"   ✅ {updated_count} lignes mises à jour dans '{table_name}'")


def migrate():
    """Exécute la migration sur toutes les tables"""
    print("🚀 Démarrage de la migration SAST (TOUTES TABLES)")
    print(f"📦 Base de données: {DATABASE_URL.split('@')[-1]}")
    
    try:
        # Connexion à PostgreSQL
        conn = psycopg2.connect(DATABASE_URL)
        cursor = conn.cursor()
        
        # Vérifier quelles tables existent
        cursor.execute("""
            SELECT table_name 
            FROM information_schema.tables 
            WHERE table_schema = 'public' 
              AND table_name IN ('api_key', 'users')
            ORDER BY table_name
        """)
        tables = [row[0] for row in cursor.fetchall()]
        print(f"\n📋 Tables trouvées: {', '.join(tables)}")
        
        # Migrer chaque table
        for table in tables:
            migrate_table(cursor, conn, table)
        
        # Vérification finale pour chaque table
        print("\n📈 Vérification des résultats...")
        
        for table in tables:
            print(f"\n📊 Table '{table}':")
            cursor.execute(f"""
                SELECT 
                    tier,
                    COUNT(*) as total,
                    COUNT(CASE WHEN allow_source_scan = TRUE THEN 1 END) as with_sast,
                    AVG(max_source_files) as avg_files,
                    AVG(max_source_size_mb) as avg_size
                FROM {table}
                GROUP BY tier
                ORDER BY 
                    CASE tier
                        WHEN 'SYSOP' THEN 1
                        WHEN 'sysop' THEN 1
                        WHEN 'ENTERPRISE' THEN 2
                        WHEN 'enterprise' THEN 2
                        WHEN 'PRO' THEN 3
                        WHEN 'pro' THEN 3
                        WHEN 'STARTER' THEN 4
                        WHEN 'starter' THEN 4
                        WHEN 'FREE' THEN 5
                        WHEN 'free' THEN 5
                        ELSE 6
                    END
            """)
            
            results = cursor.fetchall()
            
            if results:
                print("┌─────────────┬───────┬──────────┬──────────────┬─────────────┐")
                print("│ Tier        │ Total │ Avec SAST│ Moy. Fichiers│ Moy. Taille │")
                print("├─────────────┼───────┼──────────┼──────────────┼─────────────┤")
                
                for tier, total, with_sast, avg_files, avg_size in results:
                    tier_display = tier or 'NULL'
                    avg_files_val = avg_files if avg_files is not None else 0
                    avg_size_val = avg_size if avg_size is not None else 0
                    print(f"│ {tier_display:<11} │ {total:>5} │ {with_sast:>8} │ {avg_files_val:>12.0f} │ {avg_size_val:>10.0f} MB │")
                
                print("└─────────────┴───────┴──────────┴──────────────┴─────────────┘")
            else:
                print("   (Aucune donnée)")
        
        cursor.close()
        conn.close()
        
        print("\n✅ Migration terminée avec succès!")
        print("\n💡 Prochaines étapes:")
        print("   1. Tester avec: python scripts/database/test_sast_quotas.py")
        print("   2. Redémarrer le backend admin: cd admin/backend && uvicorn app.main:app --reload")
        print("   3. Vérifier via l'interface admin")
        
        return True
        
    except psycopg2.Error as e:
        print(f"\n❌ Erreur PostgreSQL: {e}")
        return False
    except Exception as e:
        print(f"\n❌ Erreur: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == '__main__':
    success = migrate()
    sys.exit(0 if success else 1)