#!/usr/bin/env python3
"""
Script de migration : Ajout des colonnes SAST dans PostgreSQL
Date: 7 novembre 2025
Usage: python scripts/database/migrate_sast_columns.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()

# 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():
    """Exécute la migration"""
    print("🚀 Démarrage de la migration SAST...")
    print(f"📦 Base de données: {DATABASE_URL.split('@')[-1]}")
    
    try:
        # Connexion à PostgreSQL
        conn = psycopg2.connect(DATABASE_URL)
        cursor = conn.cursor()
        
        # Étape 1: Ajouter les colonnes si elles n'existent pas
        print("\n📝 Étape 1: Ajout des colonnes SAST...")
        cursor.execute("""
            ALTER TABLE api_key 
                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("   ✅ Colonnes ajoutées")
        
        # Étape 2: Récupérer les licences existantes
        print("\n📊 Étape 2: Récupération des licences existantes...")
        cursor.execute("SELECT id, email, tier FROM api_key")
        licenses = cursor.fetchall()
        print(f"   📌 {len(licenses)} licences trouvées")
        
        # Étape 3: Mettre à jour chaque licence selon son tier
        print("\n🔄 Étape 3: Mise à jour des quotas par tier...")
        updated_count = 0
        
        for license_id, email, tier in licenses:
            tier_upper = tier.upper()
            quotas = TIER_QUOTAS.get(tier_upper, TIER_QUOTAS['FREE'])
            
            cursor.execute("""
                UPDATE api_key 
                SET 
                    allow_source_scan = %s,
                    max_source_files = %s,
                    max_source_size_mb = %s,
                    advanced_rules = %s,
                    updated_at = NOW()
                WHERE id = %s
            """, (
                quotas['allow_source_scan'],
                quotas['max_source_files'],
                quotas['max_source_size_mb'],
                quotas['advanced_rules'],
                license_id
            ))
            updated_count += 1
            
            print(f"   ✓ {email} ({tier_upper}): "
                  f"{quotas['max_source_files']} fichiers, "
                  f"{quotas['max_source_size_mb']} MB, "
                  f"règles avancées: {quotas['advanced_rules']}")
        
        conn.commit()
        print(f"\n   ✅ {updated_count} licences mises à jour")
        
        # Étape 4: Vérification finale
        print("\n📈 Étape 4: Vérification des résultats...")
        cursor.execute("""
            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 api_key
            GROUP BY tier
            ORDER BY 
                CASE tier
                    WHEN 'SYSOP' THEN 1
                    WHEN 'ENTERPRISE' THEN 2
                    WHEN 'PRO' THEN 3
                    WHEN 'STARTER' THEN 4
                    WHEN 'FREE' THEN 5
                    ELSE 6
                END
        """)
        
        results = cursor.fetchall()
        print("\n┌─────────────┬───────┬──────────┬──────────────┬─────────────┐")
        print("│ Tier        │ Total │ Avec SAST│ Moy. Fichiers│ Moy. Taille │")
        print("├─────────────┼───────┼──────────┼──────────────┼─────────────┤")
        
        for tier, total, with_sast, avg_files, avg_size in results:
            print(f"│ {tier:<11} │ {total:>5} │ {with_sast:>8} │ {avg_files:>12.0f} │ {avg_size:>10.0f} MB │")
        
        print("└─────────────┴───────┴──────────┴──────────────┴─────────────┘")
        
        cursor.close()
        conn.close()
        
        print("\n✅ Migration terminée avec succès!")
        print("\n💡 Prochaines étapes:")
        print("   1. Redémarrer le backend admin: cd admin/backend && uvicorn app.main:app --reload")
        print("   2. Tester la création/édition de licences via l'interface admin")
        print("   3. Vérifier que les webhooks Stripe appliquent les bons quotas")
        
        return True
        
    except psycopg2.Error as e:
        print(f"\n❌ Erreur PostgreSQL: {e}")
        return False
    except Exception as e:
        print(f"\n❌ Erreur: {e}")
        return False


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