"""
Simple migration to add predefined task columns
"""
import asyncio
import sys
from pathlib import Path

# Add backend to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from sqlalchemy import text
from app.database import engine


async def run_simple_migration():
    """Add the new columns to scheduled_jobs table"""
    async with engine.begin() as conn:
        print("🔄 Adding new columns to scheduled_jobs table...")

        # Add job_type column
        try:
            await conn.execute(text("""
                ALTER TABLE scheduled_jobs
                ADD COLUMN job_type jobtype DEFAULT 'custom_code'::jobtype NOT NULL
            """))
            print("✅ Added job_type column")
        except Exception as e:
            print(f"⚠️  job_type column: {e}")

        # Add predefined_function column
        try:
            await conn.execute(text("""
                ALTER TABLE scheduled_jobs
                ADD COLUMN predefined_function VARCHAR(100)
            """))
            print("✅ Added predefined_function column")
        except Exception as e:
            print(f"⚠️  predefined_function column: {e}")

        # Add function_params column
        try:
            await conn.execute(text("""
                ALTER TABLE scheduled_jobs
                ADD COLUMN function_params JSONB DEFAULT '{}'::jsonb
            """))
            print("✅ Added function_params column")
        except Exception as e:
            print(f"⚠️  function_params column: {e}")

        # Set default for future records to 'predefined'
        try:
            await conn.execute(text("""
                ALTER TABLE scheduled_jobs
                ALTER COLUMN job_type SET DEFAULT 'predefined'::jobtype
            """))
            print("✅ Set default job_type to 'predefined'")
        except Exception as e:
            print(f"⚠️  Setting default: {e}")

        # Add check constraint
        try:
            await conn.execute(text("""
                ALTER TABLE scheduled_jobs
                DROP CONSTRAINT IF EXISTS chk_job_type_fields
            """))
            await conn.execute(text("""
                ALTER TABLE scheduled_jobs
                ADD CONSTRAINT chk_job_type_fields
                CHECK (
                    (job_type = 'predefined' AND predefined_function IS NOT NULL)
                    OR
                    (job_type = 'custom_code' AND python_code IS NOT NULL)
                )
            """))
            print("✅ Added check constraint")
        except Exception as e:
            print(f"⚠️  Check constraint: {e}")

        print("\n📋 Final table structure:")
        result = await conn.execute(text("""
            SELECT column_name, data_type, is_nullable
            FROM information_schema.columns
            WHERE table_name = 'scheduled_jobs'
            ORDER BY ordinal_position
        """))

        rows = result.fetchall()
        for row in rows:
            print(f"  - {row[0]}: {row[1]} (nullable: {row[2]})")


if __name__ == "__main__":
    asyncio.run(run_simple_migration())