"""
Migration script to add image fields to projects and generated_content tables
Story 7.1: Generate and Insert Images into Articles
"""

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.database.session import db_manager
from sqlalchemy import text


def migrate():
    """Add image fields to projects and generated_content tables"""
    
    session = db_manager.get_session()
    
    try:
        print("Starting migration: Add image fields...")
        
        print("  Adding image_theme_prompt to projects table...")
        session.execute(text("""
            ALTER TABLE projects 
            ADD COLUMN image_theme_prompt TEXT NULL
        """))
        
        print("  Adding hero_image_url to generated_content table...")
        session.execute(text("""
            ALTER TABLE generated_content 
            ADD COLUMN hero_image_url TEXT NULL
        """))
        
        print("  Adding content_images to generated_content table...")
        session.execute(text("""
            ALTER TABLE generated_content 
            ADD COLUMN content_images JSON NULL
        """))
        
        session.commit()
        
        print("Migration completed successfully!")
        print("\nNew fields added:")
        print("  - projects.image_theme_prompt (TEXT, nullable)")
        print("  - generated_content.hero_image_url (TEXT, nullable)")
        print("  - generated_content.content_images (JSON, nullable)")
        
    except Exception as e:
        session.rollback()
        print(f"Migration failed: {e}")
        raise
    
    finally:
        session.close()


def rollback():
    """Rollback migration (remove image fields)"""
    
    session = db_manager.get_session()
    
    try:
        print("Rolling back migration: Remove image fields...")
        
        print("  Removing content_images column...")
        session.execute(text("""
            ALTER TABLE generated_content 
            DROP COLUMN content_images
        """))
        
        print("  Removing hero_image_url column...")
        session.execute(text("""
            ALTER TABLE generated_content 
            DROP COLUMN hero_image_url
        """))
        
        print("  Removing image_theme_prompt column...")
        session.execute(text("""
            ALTER TABLE projects 
            DROP COLUMN image_theme_prompt
        """))
        
        session.commit()
        
        print("Rollback completed successfully!")
        
    except Exception as e:
        session.rollback()
        print(f"Rollback failed: {e}")
        raise
    
    finally:
        session.close()


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "rollback":
        rollback()
    else:
        migrate()