import sys
import os

# Add the project root to the python path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy import text
from core.database import engine

def migrate():
    print("Starting migration: Ensuring source_id exists in all output tables...")
    
    tables = [
        "podcasts",
        "video_summaries",
        "mind_maps",
        "flashcard_sets",
        "quiz_sets",
        "reports"
    ]
    
    try:
        with engine.connect() as connection:
            for table in tables:
                print(f"Processing table: {table}...")
                
                # Check and add column/constraint
                sql = f"""
                IF NOT EXISTS (SELECT * FROM sys.columns WHERE object_id = OBJECT_ID(N'[dbo].[{table}]') AND name = 'source_id')
                BEGIN
                    ALTER TABLE [dbo].[{table}] ADD [source_id] INT NULL;
                    ALTER TABLE [dbo].[{table}] ADD CONSTRAINT [FK_{table}_sources] 
                    FOREIGN KEY ([source_id]) REFERENCES [dbo].[sources] ([id]);
                END
                """
                connection.execute(text(sql))
            
            connection.commit()
            print("Migration completed successfully!")
    except Exception as e:
        print(f"Error during migration: {e}")

if __name__ == "__main__":
    migrate()