"""
Migrate data from SQLite to PostgreSQL.

This script:
1. Reads all data from the existing SQLite database
2. Reads asset files from the output/ directory
3. Inserts everything into PostgreSQL with BYTEA for binary assets

Usage:
    python -m scripts.migrate_sqlite_to_postgres
"""

import asyncio
import json
import os
import sys
from pathlib import Path
from datetime import datetime, timezone

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

import aiosqlite
import asyncpg
from dotenv import load_dotenv

load_dotenv()


def get_database_url() -> str:
    """Get PostgreSQL database URL from environment."""
    url = os.getenv("DATABASE_URL")
    if not url:
        url = "postgresql://video_automation:video_automation@localhost:5432/video_automation"
    return url


def get_sqlite_path() -> Path:
    """Get SQLite database path."""
    return project_root / "data" / "pipeline.db"


async def migrate_projects(sqlite_conn: aiosqlite.Connection, pg_conn: asyncpg.Connection) -> dict[str, str]:
    """Migrate projects table. Returns mapping of old_id -> new_id."""
    print("  Migrating projects...")
    
    sqlite_conn.row_factory = aiosqlite.Row
    cursor = await sqlite_conn.execute("SELECT * FROM projects")
    rows = await cursor.fetchall()
    
    id_mapping = {}
    
    for row in rows:
        old_id = row["id"]
        
        # Check if already exists
        existing = await pg_conn.fetchrow(
            "SELECT id FROM projects WHERE id = $1::uuid", old_id
        )
        if existing:
            id_mapping[old_id] = old_id
            continue
        
        try:
            await pg_conn.execute(
                """INSERT INTO projects (id, status, mode, platform, character,
                                         story_preview, title, created_at, result_json)
                   VALUES ($1::uuid, $2, $3, $4, $5, $6, $7, $8, $9)
                   ON CONFLICT (id) DO NOTHING""",
                old_id,
                row["status"],
                row["mode"],
                row["platform"],
                row["character"],
                row["story_preview"],
                row["title"],
                datetime.fromisoformat(row["created_at"]) if row["created_at"] else datetime.now(timezone.utc),
                row["result_json"],
            )
            id_mapping[old_id] = old_id
        except Exception as e:
            print(f"    Warning: Failed to migrate project {old_id}: {e}")
    
    print(f"    ✓ Migrated {len(id_mapping)} projects")
    return id_mapping


async def migrate_characters(sqlite_conn: aiosqlite.Connection, pg_conn: asyncpg.Connection) -> int:
    """Migrate characters table."""
    print("  Migrating characters...")
    
    sqlite_conn.row_factory = aiosqlite.Row
    cursor = await sqlite_conn.execute("SELECT * FROM characters")
    rows = await cursor.fetchall()
    
    count = 0
    for row in rows:
        char_id = row["id"]
        
        # Check if already exists
        existing = await pg_conn.fetchrow(
            "SELECT id FROM characters WHERE id = $1::uuid OR name = $2",
            char_id, row["name"]
        )
        if existing:
            continue
        
        try:
            await pg_conn.execute(
                """INSERT INTO characters
                   (id, name, physical_description, style_keywords, negative_keywords,
                    hard_prohibitions, camera_settings, lighting_settings, pose_rule,
                    environment_rule, clothing_rule, anatomical_highlight_rules,
                    thumbnail_path, is_default, created_at, updated_at)
                   VALUES ($1::uuid, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
                   ON CONFLICT (id) DO NOTHING""",
                char_id,
                row["name"],
                row["physical_description"],
                row["style_keywords"],
                row["negative_keywords"],
                row["hard_prohibitions"],
                row["camera_settings"],
                row["lighting_settings"],
                row["pose_rule"],
                row["environment_rule"],
                row["clothing_rule"],
                row["anatomical_highlight_rules"],
                row["thumbnail_path"],
                bool(row["is_default"]),
                datetime.fromisoformat(row["created_at"]) if row["created_at"] else datetime.now(timezone.utc),
                datetime.fromisoformat(row["updated_at"]) if row["updated_at"] else datetime.now(timezone.utc),
            )
            count += 1
        except Exception as e:
            print(f"    Warning: Failed to migrate character {row['name']}: {e}")
    
    print(f"    ✓ Migrated {count} characters")
    return count


async def migrate_images(
    sqlite_conn: aiosqlite.Connection, 
    pg_conn: asyncpg.Connection,
    project_ids: dict[str, str]
) -> int:
    """Migrate images table with BYTEA data from files."""
    print("  Migrating images...")
    
    sqlite_conn.row_factory = aiosqlite.Row
    cursor = await sqlite_conn.execute("SELECT * FROM images")
    rows = await cursor.fetchall()
    
    count = 0
    for row in rows:
        project_id = row["project_id"]
        if project_id not in project_ids:
            continue
        
        # Check if already exists
        existing = await pg_conn.fetchrow(
            "SELECT id FROM images WHERE project_id = $1::uuid AND scene_number = $2",
            project_id, row["scene_number"]
        )
        if existing:
            continue
        
        # Try to read image file
        image_data = None
        file_path = row["file_path"]
        if file_path and Path(file_path).exists():
            try:
                image_data = Path(file_path).read_bytes()
            except Exception as e:
                print(f"    Warning: Could not read image file {file_path}: {e}")
        
        try:
            await pg_conn.execute(
                """INSERT INTO images (project_id, scene_number, image_prompt,
                                       file_path, image_data, status, created_at, updated_at)
                   VALUES ($1::uuid, $2, $3, $4, $5, $6, $7, $8)
                   ON CONFLICT (project_id, scene_number) DO NOTHING""",
                project_id,
                row["scene_number"],
                row["image_prompt"],
                file_path,
                image_data,
                row["status"],
                datetime.fromisoformat(row["created_at"]) if row["created_at"] else datetime.now(timezone.utc),
                datetime.fromisoformat(row["updated_at"]) if row["updated_at"] else datetime.now(timezone.utc),
            )
            count += 1
        except Exception as e:
            print(f"    Warning: Failed to migrate image {project_id}/{row['scene_number']}: {e}")
    
    print(f"    ✓ Migrated {count} images")
    return count


async def migrate_videos(
    sqlite_conn: aiosqlite.Connection, 
    pg_conn: asyncpg.Connection,
    project_ids: dict[str, str]
) -> int:
    """Migrate videos table with BYTEA data from files."""
    print("  Migrating videos...")
    
    sqlite_conn.row_factory = aiosqlite.Row
    cursor = await sqlite_conn.execute("SELECT * FROM videos")
    rows = await cursor.fetchall()
    
    count = 0
    for row in rows:
        project_id = row["project_id"]
        if project_id not in project_ids:
            continue
        
        # Check if already exists
        existing = await pg_conn.fetchrow(
            "SELECT id FROM videos WHERE project_id = $1::uuid AND scene_number = $2",
            project_id, row["scene_number"]
        )
        if existing:
            continue
        
        # Try to read video file
        video_data = None
        file_path = row["file_path"]
        if file_path and Path(file_path).exists():
            try:
                video_data = Path(file_path).read_bytes()
            except Exception as e:
                print(f"    Warning: Could not read video file {file_path}: {e}")
        
        try:
            await pg_conn.execute(
                """INSERT INTO videos (project_id, scene_number, video_prompt,
                                       file_path, video_data, status, created_at, updated_at)
                   VALUES ($1::uuid, $2, $3, $4, $5, $6, $7, $8)
                   ON CONFLICT (project_id, scene_number) DO NOTHING""",
                project_id,
                row["scene_number"],
                row["video_prompt"],
                file_path,
                video_data,
                row["status"],
                datetime.fromisoformat(row["created_at"]) if row["created_at"] else datetime.now(timezone.utc),
                datetime.fromisoformat(row["updated_at"]) if row["updated_at"] else datetime.now(timezone.utc),
            )
            count += 1
        except Exception as e:
            print(f"    Warning: Failed to migrate video {project_id}/{row['scene_number']}: {e}")
    
    print(f"    ✓ Migrated {count} videos")
    return count


async def import_output_directory(pg_conn: asyncpg.Connection) -> int:
    """Import assets from output/ directory that aren't in SQLite."""
    print("  Scanning output directory for additional assets...")
    
    output_dir = project_root / "output"
    if not output_dir.exists():
        print("    No output directory found")
        return 0
    
    count = 0
    
    for project_dir in output_dir.iterdir():
        if not project_dir.is_dir():
            continue
        
        project_id = project_dir.name
        
        # Skip non-UUID directories
        if len(project_id) != 36 or project_id.count('-') != 4:
            continue
        
        # Check if project exists in DB
        existing = await pg_conn.fetchrow(
            "SELECT id FROM projects WHERE id = $1::uuid", project_id
        )
        
        if not existing:
            # Create project from production_plan.json if exists
            plan_file = project_dir / "production_plan.json"
            if plan_file.exists():
                try:
                    plan_data = json.loads(plan_file.read_text())
                    await pg_conn.execute(
                        """INSERT INTO projects (id, status, mode, platform, character,
                                                 story_preview, title, created_at, result_json)
                           VALUES ($1::uuid, $2, $3, $4, $5, $6, $7, $8, $9)
                           ON CONFLICT (id) DO NOTHING""",
                        project_id,
                        "completed",
                        "generate",
                        plan_data.get("target_platform", "tiktok"),
                        plan_data.get("character_template_name", ""),
                        plan_data.get("summary", "")[:500],
                        plan_data.get("title", "Untitled"),
                        datetime.now(timezone.utc),
                        json.dumps(plan_data),
                    )
                    count += 1
                except Exception as e:
                    print(f"    Warning: Could not import project {project_id}: {e}")
                    continue
        
        # Import assets
        assets_dir = project_dir / "assets"
        if assets_dir.exists():
            for asset_file in assets_dir.iterdir():
                if not asset_file.is_file():
                    continue
                
                # Parse filename: frame_01.png, frame_01.mp4, etc.
                name = asset_file.stem
                if not name.startswith("frame_"):
                    continue
                
                try:
                    scene_number = int(name.split("_")[1])
                except (IndexError, ValueError):
                    continue
                
                suffix = asset_file.suffix.lower()
                
                if suffix in (".png", ".jpg", ".jpeg", ".webp"):
                    # Image
                    existing = await pg_conn.fetchrow(
                        "SELECT id FROM images WHERE project_id = $1::uuid AND scene_number = $2",
                        project_id, scene_number
                    )
                    if not existing:
                        try:
                            image_data = asset_file.read_bytes()
                            mime_type = {
                                ".png": "image/png",
                                ".jpg": "image/jpeg",
                                ".jpeg": "image/jpeg",
                                ".webp": "image/webp",
                            }.get(suffix, "image/png")
                            
                            await pg_conn.execute(
                                """INSERT INTO images (project_id, scene_number, file_path,
                                                       image_data, mime_type, file_size_bytes,
                                                       status, created_at, updated_at)
                                   VALUES ($1::uuid, $2, $3, $4, $5, $6, $7, $8, $9)
                                   ON CONFLICT (project_id, scene_number) DO UPDATE
                                   SET image_data = EXCLUDED.image_data,
                                       mime_type = EXCLUDED.mime_type,
                                       file_size_bytes = EXCLUDED.file_size_bytes""",
                                project_id, scene_number, str(asset_file),
                                image_data, mime_type, len(image_data),
                                "generated", datetime.now(timezone.utc), datetime.now(timezone.utc),
                            )
                            count += 1
                        except Exception as e:
                            print(f"    Warning: Could not import image {asset_file}: {e}")
                
                elif suffix in (".mp4", ".webm", ".mov"):
                    # Video
                    existing = await pg_conn.fetchrow(
                        "SELECT id FROM videos WHERE project_id = $1::uuid AND scene_number = $2",
                        project_id, scene_number
                    )
                    if not existing:
                        try:
                            video_data = asset_file.read_bytes()
                            mime_type = {
                                ".mp4": "video/mp4",
                                ".webm": "video/webm",
                                ".mov": "video/quicktime",
                            }.get(suffix, "video/mp4")
                            
                            await pg_conn.execute(
                                """INSERT INTO videos (project_id, scene_number, file_path,
                                                       video_data, mime_type, file_size_bytes,
                                                       status, created_at, updated_at)
                                   VALUES ($1::uuid, $2, $3, $4, $5, $6, $7, $8, $9)
                                   ON CONFLICT (project_id, scene_number) DO UPDATE
                                   SET video_data = EXCLUDED.video_data,
                                       mime_type = EXCLUDED.mime_type,
                                       file_size_bytes = EXCLUDED.file_size_bytes""",
                                project_id, scene_number, str(asset_file),
                                video_data, mime_type, len(video_data),
                                "generated", datetime.now(timezone.utc), datetime.now(timezone.utc),
                            )
                            count += 1
                        except Exception as e:
                            print(f"    Warning: Could not import video {asset_file}: {e}")
    
    print(f"    ✓ Imported {count} additional assets from output/")
    return count


async def run_migration():
    """Run the full migration."""
    print("=" * 60)
    print("SQLite to PostgreSQL Migration")
    print("=" * 60)
    
    sqlite_path = get_sqlite_path()
    pg_url = get_database_url()
    
    print(f"\nSource: {sqlite_path}")
    print(f"Target: {pg_url.split('@')[1] if '@' in pg_url else pg_url}")
    print()
    
    # Connect to PostgreSQL
    print("Connecting to PostgreSQL...")
    try:
        pg_conn = await asyncpg.connect(pg_url)
    except Exception as e:
        print(f"Failed to connect to PostgreSQL: {e}")
        print("Make sure PostgreSQL is running (use 'make up' or 'docker compose up -d db')")
        sys.exit(1)
    
    try:
        # Check if SQLite exists
        if sqlite_path.exists():
            print("Connecting to SQLite...")
            sqlite_conn = await aiosqlite.connect(str(sqlite_path))
            
            try:
                # Migrate data
                print("\nMigrating data from SQLite...")
                project_ids = await migrate_projects(sqlite_conn, pg_conn)
                await migrate_characters(sqlite_conn, pg_conn)
                await migrate_images(sqlite_conn, pg_conn, project_ids)
                await migrate_videos(sqlite_conn, pg_conn, project_ids)
            finally:
                await sqlite_conn.close()
        else:
            print(f"No SQLite database found at {sqlite_path}")
        
        # Import from output directory
        print("\nImporting from output directory...")
        await import_output_directory(pg_conn)
        
        print("\n" + "=" * 60)
        print("Migration complete!")
        print("=" * 60)
        
    finally:
        await pg_conn.close()


def main():
    """Entry point."""
    asyncio.run(run_migration())


if __name__ == "__main__":
    main()