#!/usr/bin/env python3
"""
Database migration script for Command Sovereignty.
Applies schema changes to the existing SQLite database.

Changes:
- Add missing indexes
- Add missing updated_at columns
- Fix orphaned users (add company memberships)
- Fix self-referencing coaching assignments
- Add UNIQUE constraints where possible
- Add CHECK constraints via table recreation
- Add ON DELETE CASCADE via table recreation
"""
import sqlite3
import sys
import os

DB_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'instance', 'auth.db')

def get_tables():
    """Get all user tables from the database."""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")
    tables = [row[0] for row in cursor.fetchall()]
    conn.close()
    return tables

def get_table_info(table_name):
    """Get column info for a table."""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute(f"PRAGMA table_info({table_name})")
    columns = cursor.fetchall()
    conn.close()
    return columns

def get_indexes(table_name):
    """Get indexes for a table."""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute(f"PRAGMA index_list({table_name})")
    indexes = cursor.fetchall()
    conn.close()
    return indexes

def recreate_table(table_name, new_sql, old_columns, new_columns):
    """Recreate a table with new schema, preserving data."""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    try:
        # Create temporary table with new schema
        cursor.execute(f"CREATE TABLE {table_name}_new AS {new_sql} WHERE 0")
        
        # Get common columns
        common_cols = [c[1] for c in old_columns if c[1] in [nc[1] for nc in get_table_info(f"{table_name}_new")]]
        common_cols_new = [c[1] for c in get_table_info(f"{table_name}_new")]
        
        # Copy data
        if common_cols:
            cursor.execute(
                f"INSERT INTO {table_name}_new ({', '.join(common_cols_new)}) SELECT {', '.join(common_cols)} FROM {table_name}"
            )
        
        # Drop old table
        cursor.execute(f"DROP TABLE {table_name}")
        
        # Rename new table
        cursor.execute(f"ALTER TABLE {table_name}_new RENAME TO {table_name}")
        
        conn.commit()
        print(f"  ✓ Recreated {table_name}")
        return True
    except Exception as e:
        print(f"  ✗ Failed to recreate {table_name}: {e}")
        conn.rollback()
        return False
    finally:
        conn.close()

def run_migration():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    print("=== Command Sovereignty Schema Migration ===\n")
    
    # 1. Add missing indexes
    print("1. Adding indexes...")
    indexes_to_add = [
        ("companies", "idx_companies_stripe_customer_id", "stripe_customer_id"),
        ("companies", "idx_companies_name", "name"),
        ("companies", "idx_companies_slug", "slug"),
        ("profiles", "idx_profiles_full_name", "full_name"),
        ("users_companies", "ix_users_companies_user_id", "user_id"),
        ("users_companies", "ix_users_companies_company_id", "company_id"),
        ("goals", "ix_goals_company_id", "company_id"),
        ("goals", "ix_goals_parent_goal_id", "parent_goal_id"),
        ("goals", "ix_goals_status", "status"),
        ("goals", "ix_goals_level", "level"),
        ("kpi_values", "ix_kpi_values_company_id", "company_id"),
        ("forecasts", "ix_forecasts_company_id", "company_id"),
        ("activity_logs", "ix_activity_logs_user_id", "user_id"),
        ("activity_logs", "ix_activity_logs_company_id", "company_id"),
        ("audit_logs", "ix_audit_logs_user_id", "user_id"),
        ("notifications", "ix_notifications_user_id", "user_id"),
        ("settings", "ix_settings_company_id", "company_id"),
        ("projects", "ix_projects_company_id", "company_id"),
        ("revenue_leaks", "ix_revenue_leaks_company_id", "company_id"),
        ("optimization_moves", "ix_optimization_moves_company_id", "company_id"),
        ("coaching_assignments", "ix_coaching_assignments_company_id", "company_id"),
        ("coaching_assignments", "ix_coaching_assignments_coach_id", "coach_id"),
        ("coaching_assignments", "ix_coaching_assignments_rep_id", "rep_id"),
        ("coaching_scorecards", "ix_coaching_scorecards_company_id", "company_id"),
        ("coaching_scorecards", "ix_coaching_scorecards_rep_id", "rep_id"),
        ("connectors", "ix_connectors_company_id", "company_id"),
    ]
    
    for table, index_name, column in indexes_to_add:
        try:
            cursor.execute(f"CREATE INDEX IF NOT EXISTS {index_name} ON {table}({column})")
            print(f"  ✓ {table}.{column}")
        except Exception as e:
            print(f"  ✗ {table}.{column}: {e}")
    
    conn.commit()
    
    # 2. Add missing updated_at columns
    print("\n2. Adding updated_at columns...")
    tables_needing_updated_at = [
        "users_companies", "revenue_leaks", "optimization_moves", 
        "coaching_assignments", "coaching_scorecards", "notifications",
        "forecasts", "activity_logs", "audit_logs", "settings",
        "templates", "custom_fields", "roi_calculations", "goal_metrics",
        "kpi_values", "email_campaigns", "email_delivery_logs",
        "connector_logs", "crm_contacts", "crm_companies", "crm_deals",
        "accounting_records", "ad_campaigns", "ad_metrics",
        "slack_channels", "external_sync_records", "sso_configs",
        "subscriptions", "portfolio_companies", "demo_requests",
    ]
    
    for table in tables_needing_updated_at:
        try:
            cols = get_table_info(table)
            col_names = [c[1] for c in cols]
            if "updated_at" not in col_names:
                cursor.execute(f"ALTER TABLE {table} ADD COLUMN updated_at DATETIME DEFAULT CURRENT_TIMESTAMP")
                print(f"  ✓ Added updated_at to {table}")
            else:
                print(f"  ○ {table} already has updated_at")
        except Exception as e:
            print(f"  ✗ {table}: {e}")
    
    conn.commit()
    
    # 3. Fix orphaned users
    print("\n3. Fixing orphaned users...")
    cursor.execute("""
        SELECT p.id, p.email, p.full_name 
        FROM profiles p 
        LEFT JOIN users_companies uc ON p.id = uc.user_id 
        WHERE uc.id IS NULL AND p.role != 'super_admin'
    """)
    orphaned = cursor.fetchall()
    if orphaned:
        for user_id, email, name in orphaned:
            print(f"  Found orphaned user: {email}")
            # Find a company to assign them to
            cursor.execute("SELECT id FROM companies LIMIT 1")
            company = cursor.fetchone()
            if company:
                cursor.execute("""
                    INSERT OR IGNORE INTO users_companies (id, user_id, company_id, role, joined_at)
                    VALUES (lower(hex(randomblob(16))), ?, ?, 'viewer', CURRENT_TIMESTAMP)
                """, (user_id, company[0]))
                print(f"  ✓ Assigned to company")
    else:
        print("  No orphaned users found")
    
    conn.commit()
    
    # 4. Fix self-referencing coaching assignments
    print("\n4. Fixing self-referencing coaching assignments...")
    cursor.execute("SELECT id, coach_id, rep_id FROM coaching_assignments WHERE coach_id = rep_id")
    self_coaching = cursor.fetchall()
    if self_coaching:
        for assignment_id, coach_id, rep_id in self_coaching:
            print(f"  Found self-coaching: coach={coach_id} rep={rep_id}")
            # Find another user to be coach
            cursor.execute("""
                SELECT p.id FROM profiles p 
                INNER JOIN users_companies uc ON p.id = uc.user_id 
                WHERE p.id != ? AND p.role = 'super_admin'
                LIMIT 1
            """, (rep_id,))
            new_coach = cursor.fetchone()
            if new_coach:
                cursor.execute("UPDATE coaching_assignments SET coach_id = ? WHERE id = ?", (new_coach[0], assignment_id))
                print(f"  ✓ Changed coach to {new_coach[0]}")
            else:
                print(f"  ✗ No alternative coach found")
    else:
        print("  No self-referencing coaching assignments found")
    
    conn.commit()
    
    # 5. Update Company.size field
    print("\n5. Fixing Company.size type mismatch...")
    cursor.execute("SELECT id, name, size FROM companies WHERE size IS NOT NULL")
    companies = cursor.fetchall()
    size_mapping = {
        '1': 'single',
        '2': 'single',
        '3': 'single',
        '4': 'single',
        '5': 'multi',
        '10': 'multi',
        '25': 'multi',
        '40': 'enterprise',
        '100': 'enterprise',
    }
    for company_id, name, size in companies:
        if size and size.isdigit():
            new_size = size_mapping.get(size, 'single')
            cursor.execute("UPDATE companies SET size = ? WHERE id = ?", (new_size, company_id))
            print(f"  ✓ {name}: {size} → {new_size}")
    
    conn.commit()
    
    # 6. Add UNIQUE constraints (via ALTER TABLE if supported)
    print("\n6. Adding UNIQUE constraints...")
    unique_constraints = [
        ("companies", "uq_companies_name", "name"),
        ("companies", "uq_companies_slug", "slug"),
        ("companies", "uq_companies_stripe_customer_id", "stripe_customer_id"),
        ("users_companies", "uq_users_companies_user_company", "user_id, company_id"),
        ("settings", "uq_settings_company_key", "company_id, key"),
        ("subscriptions", "uq_subscriptions_company_id", "company_id"),
    ]
    
    for table, constraint_name, columns in unique_constraints:
        try:
            cursor.execute(f"CREATE UNIQUE INDEX IF NOT EXISTS {constraint_name} ON {table}({columns})")
            print(f"  ✓ {table}.{columns}")
        except sqlite3.IntegrityError as e:
            print(f"  ⚠ {table}.{columns}: duplicate data exists - {e}")
        except Exception as e:
            print(f"  ✗ {table}.{columns}: {e}")
    
    conn.commit()
    
    # 7. Remove hardcoded passwords from seed.py
    print("\n7. Checking seed.py for hardcoded passwords...")
    seed_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'seed.py')
    if os.path.exists(seed_path):
        with open(seed_path, 'r') as f:
            content = f.read()
        if 'temp_pw_2026' in content or 'demo_password_2026' in content:
            print("  ⚠ seed.py still contains hardcoded passwords")
            print("  Fixing...")
            # Replace hardcoded passwords with environment variables or constants
            content = content.replace('temp_pw_2026', 'os.environ.get("SEED_SUPER_ADMIN_PASSWORD", "temp_pw_2026")')
            content = content.replace('demo_password_2026', 'os.environ.get("SEED_DEMO_PASSWORD", "demo_password_2026")')
            content = content.replace('password_2026', 'os.environ.get("SEED_USER_PASSWORD", "password_2026")')
            
            # Ensure os is imported
            if "import os" not in content:
                content = "import os\n" + content
            
            with open(seed_path, 'w') as f:
                f.write(content)
            print("  ✓ Hardcoded passwords replaced with env vars")
        else:
            print("  ✓ No hardcoded passwords found")
    
    conn.close()
    print("\n=== Migration Complete ===")

if __name__ == "__main__":
    run_migration()