#!/usr/bin/env python3
"""
Migrate data from SQLite to PostgreSQL database.

This script:
1. Creates PostgreSQL databases if they don't exist
2. Creates all tables in PostgreSQL
3. Migrates all data from SQLite to PostgreSQL
"""

from __future__ import annotations

import sys
import os
from pathlib import Path

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import sessionmaker
from backend.db.models import Base
from backend.db.session import get_database_url
import sqlite3

def create_postgresql_databases():
    """Create PostgreSQL databases if they don't exist."""
    from urllib.parse import urlparse
    
    # Get database URLs from environment
    db_url = os.getenv("DATABASE_URL", "postgresql+psycopg2://punimtag:punimtag_password@localhost:5432/punimtag")
    auth_db_url = os.getenv("DATABASE_URL_AUTH", "postgresql://punimtag:punimtag_password@localhost:5432/punimtag_auth")
    
    # Parse URLs
    main_parsed = urlparse(db_url.replace("postgresql+psycopg2://", "postgresql://"))
    auth_parsed = urlparse(auth_db_url.replace("postgresql+psycopg2://", "postgresql://"))
    
    main_db_name = main_parsed.path.lstrip("/")
    auth_db_name = auth_parsed.path.lstrip("/")
    
    # Connect to postgres database to create other databases
    postgres_url = f"postgresql://{main_parsed.username}:{main_parsed.password}@{main_parsed.hostname}:{main_parsed.port or 5432}/postgres"
    
    try:
        engine = create_engine(postgres_url)
        with engine.connect() as conn:
            # Check if databases exist
            result = conn.execute(text("SELECT 1 FROM pg_database WHERE datname = :name"), {"name": main_db_name})
            if not result.fetchone():
                conn.execute(text("COMMIT"))  # End any transaction
                conn.execute(text(f'CREATE DATABASE "{main_db_name}"'))
                print(f"✅ Created database: {main_db_name}")
            else:
                print(f"✅ Database already exists: {main_db_name}")
            
            result = conn.execute(text("SELECT 1 FROM pg_database WHERE datname = :name"), {"name": auth_db_name})
            if not result.fetchone():
                conn.execute(text("COMMIT"))
                conn.execute(text(f'CREATE DATABASE "{auth_db_name}"'))
                print(f"✅ Created database: {auth_db_name}")
            else:
                print(f"✅ Database already exists: {auth_db_name}")
    except Exception as e:
        print(f"⚠️  Error creating databases: {e}")
        print("   Make sure PostgreSQL is running and credentials are correct")

def migrate_data():
    """Migrate data from SQLite to PostgreSQL."""
    print("=" * 80)
    print("MIGRATING DATA FROM SQLITE TO POSTGRESQL")
    print("=" * 80)
    
    # Get database URLs
    sqlite_url = "sqlite:///data/punimtag.db"
    postgres_url = os.getenv("DATABASE_URL", "postgresql+psycopg2://punimtag:punimtag_password@localhost:5432/punimtag")
    
    if not postgres_url.startswith("postgresql"):
        print("❌ DATABASE_URL is not set to PostgreSQL")
        print("   Set DATABASE_URL in .env file to PostgreSQL connection string")
        return False
    
    # Connect to both databases
    sqlite_engine = create_engine(sqlite_url)
    postgres_engine = create_engine(postgres_url)
    
    # Create tables in PostgreSQL
    print("\n📋 Creating tables in PostgreSQL...")
    Base.metadata.create_all(bind=postgres_engine)
    print("✅ Tables created")
    
    # Get table names
    inspector = inspect(sqlite_engine)
    all_tables = inspector.get_table_names()
    
    # Exclude system tables
    all_tables = [t for t in all_tables if not t.startswith("sqlite_")]
    
    # Define migration order (respecting foreign key constraints)
    # Tables with no dependencies first, then dependent tables
    migration_order = [
        "alembic_version",  # Migration tracking (optional)
        "photos",           # Base table
        "people",           # Base table
        "tags",             # Base table
        "users",            # Base table
        "faces",            # Depends on photos, people, users
        "person_encodings", # Depends on people, faces
        "phototaglinkage",  # Depends on photos, tags
        "photo_favorites",  # Depends on photos
        "photo_person_linkage", # Depends on photos, people, users
        "role_permissions", # Base table
    ]
    
    # Filter to only tables that exist
    tables = [t for t in migration_order if t in all_tables]
    # Add any remaining tables not in the order list
    for t in all_tables:
        if t not in tables:
            tables.append(t)
    
    print(f"\n📊 Found {len(tables)} tables to migrate: {', '.join(tables)}")
    
    # Boolean columns mapping (SQLite stores as integer, PostgreSQL needs boolean)
    boolean_columns = {
        "photos": ["processed"],
        "faces": ["is_primary_encoding", "excluded"],
        "users": ["is_active", "is_admin", "password_change_required"],
        "role_permissions": ["allowed"],
    }
    
    # Columns that might be missing in SQLite but required in PostgreSQL
    # Map: table_name -> {column: default_value}
    default_values = {
        "photos": {"file_hash": "migrated"},  # file_hash might be missing in old SQLite
    }
    
    # Migrate each table
    with sqlite_engine.connect() as sqlite_conn, postgres_engine.connect() as postgres_conn:
        for table in tables:
            print(f"\n🔄 Migrating table: {table}")
            
            # Get row count
            count_result = sqlite_conn.execute(text(f"SELECT COUNT(*) FROM {table}"))
            row_count = count_result.scalar()
            
            if row_count == 0:
                print(f"   ⏭️  Table is empty, skipping")
                continue
            
            print(f"   📦 {row_count} rows to migrate")
            
            # Check if table already has data in PostgreSQL
            try:
                pg_count_result = postgres_conn.execute(text(f'SELECT COUNT(*) FROM "{table}"'))
                pg_count = pg_count_result.scalar()
                if pg_count > 0:
                    print(f"   ⚠️  Table already has {pg_count} rows in PostgreSQL")
                    # Auto-truncate for non-interactive mode, or ask in interactive
                    print(f"   🗑️  Truncating existing data...")
                    postgres_conn.execute(text(f'TRUNCATE TABLE "{table}" CASCADE'))
                    postgres_conn.commit()
            except Exception as e:
                # Table might not exist yet, that's OK
                pass
            
            # Get column names and types from SQLite
            columns_result = sqlite_conn.execute(text(f"PRAGMA table_info({table})"))
            column_info = columns_result.fetchall()
            sqlite_columns = [row[1] for row in column_info]
            
            # Get PostgreSQL column names
            pg_inspector = inspect(postgres_engine)
            pg_columns_info = pg_inspector.get_columns(table)
            pg_columns = [col['name'] for col in pg_columns_info]
            
            # Use PostgreSQL columns (they're the source of truth)
            columns = pg_columns
            
            # Get boolean columns for this table
            table_bool_cols = boolean_columns.get(table, [])
            
            # Get default values for missing columns
            table_defaults = default_values.get(table, {})
            
            # Build SELECT statement for SQLite (only select columns that exist)
            select_cols = [col for col in columns if col in sqlite_columns]
            select_sql = f"SELECT {', '.join(select_cols)} FROM {table}"
            
            # Fetch all data
            data_result = sqlite_conn.execute(text(select_sql))
            rows = data_result.fetchall()
            
            # Insert into PostgreSQL
            inserted = 0
            for row in rows:
                try:
                    # Build insert statement with boolean conversion
                    values = {}
                    for i, col in enumerate(select_cols):
                        val = row[i]
                        # Convert integer booleans to Python booleans for PostgreSQL
                        if col in table_bool_cols:
                            val = bool(val) if val is not None else None
                        values[col] = val
                    
                    # Add default values for missing columns
                    for col, default_val in table_defaults.items():
                        if col not in values and col in columns:
                            values[col] = default_val
                    
                    # Only insert columns we have values for (that exist in PostgreSQL)
                    insert_cols = [col for col in columns if col in values]
                    cols_str = ', '.join([f'"{c}"' for c in insert_cols])
                    placeholders = ', '.join([f':{c}' for c in insert_cols])
                    insert_sql = f'INSERT INTO "{table}" ({cols_str}) VALUES ({placeholders})'
                    
                    postgres_conn.execute(text(insert_sql), values)
                    inserted += 1
                    
                    if inserted % 100 == 0:
                        postgres_conn.commit()
                        print(f"   ✅ Inserted {inserted}/{row_count} rows...", end='\r')
                
                except Exception as e:
                    print(f"\n   ❌ Error inserting row: {e}")
                    print(f"   Row data: {dict(zip(columns, row))}")
                    postgres_conn.rollback()
                    break
            
            postgres_conn.commit()
            print(f"   ✅ Migrated {inserted}/{row_count} rows from {table}")
    
    print("\n" + "=" * 80)
    print("✅ MIGRATION COMPLETE")
    print("=" * 80)
    print("\nNext steps:")
    print("1. Update .env file to use PostgreSQL:")
    print("   DATABASE_URL=postgresql+psycopg2://punimtag:punimtag_password@localhost:5432/punimtag")
    print("2. Restart the backend API")
    print("3. Restart the viewer frontend")
    print("4. Verify data in viewer frontend")
    
    return True

if __name__ == "__main__":
    print("🔧 SQLite to PostgreSQL Migration Tool\n")
    
    # Check if SQLite database exists
    sqlite_path = project_root / "data" / "punimtag.db"
    if not sqlite_path.exists():
        print(f"❌ SQLite database not found: {sqlite_path}")
        sys.exit(1)
    
    print(f"✅ Found SQLite database: {sqlite_path}")
    
    # Create PostgreSQL databases
    print("\n📦 Creating PostgreSQL databases...")
    create_postgresql_databases()
    
    # Migrate data
    print("\n")
    migrate_data()

