"""
Migration script to add batch_number and branch_id to serial_numbers table
"""
import sys
from pathlib import Path

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

from sqlalchemy import text
from app.core.database import engine


def main():
    print("Adding batch_number and branch_id to serial_numbers...")

    with engine.connect() as conn:
        # Check if columns already exist
        result = conn.execute(text("""
            SELECT column_name
            FROM information_schema.columns
            WHERE table_name='serial_numbers'
            AND column_name IN ('batch_number', 'branch_id')
        """))
        existing_cols = [row[0] for row in result]

        if 'batch_number' in existing_cols and 'branch_id' in existing_cols:
            print("✓ Columns already exist. Nothing to do.")
            return

        # Add batch_number column
        if 'batch_number' not in existing_cols:
            print("Adding batch_number column...")
            conn.execute(text("""
                ALTER TABLE serial_numbers
                ADD COLUMN batch_number VARCHAR NULL
            """))
            conn.commit()
            print("✓ Added batch_number column")

        # Add branch_id column
        if 'branch_id' not in existing_cols:
            print("Adding branch_id column...")
            conn.execute(text("""
                ALTER TABLE serial_numbers
                ADD COLUMN branch_id VARCHAR NULL
            """))
            conn.commit()
            print("✓ Added branch_id column")

        # Add foreign key constraint for branch_id
        print("Adding foreign key constraint...")
        try:
            conn.execute(text("""
                ALTER TABLE serial_numbers
                ADD CONSTRAINT fk_serial_numbers_branch_id
                FOREIGN KEY (branch_id) REFERENCES branches(id)
            """))
            conn.commit()
            print("✓ Added foreign key constraint")
        except Exception as e:
            if "already exists" in str(e):
                print("✓ Foreign key constraint already exists")
            else:
                raise

        # Create indexes
        print("Creating indexes...")
        indexes = [
            ("ix_serial_numbers_batch_number", "CREATE INDEX IF NOT EXISTS ix_serial_numbers_batch_number ON serial_numbers(batch_number)"),
            ("ix_serial_numbers_branch_id", "CREATE INDEX IF NOT EXISTS ix_serial_numbers_branch_id ON serial_numbers(branch_id)"),
            ("ix_serial_numbers_branch_status", "CREATE INDEX IF NOT EXISTS ix_serial_numbers_branch_status ON serial_numbers(branch_id, status)"),
            ("ix_serial_numbers_batch_status", "CREATE INDEX IF NOT EXISTS ix_serial_numbers_batch_status ON serial_numbers(batch_number, status)"),
            ("ix_serial_numbers_customer_date", "CREATE INDEX IF NOT EXISTS ix_serial_numbers_customer_date ON serial_numbers(customer_id, sold_at)"),
        ]

        for idx_name, idx_sql in indexes:
            try:
                conn.execute(text(idx_sql))
                conn.commit()
                print(f"✓ Created index: {idx_name}")
            except Exception as e:
                if "already exists" in str(e):
                    print(f"✓ Index already exists: {idx_name}")
                else:
                    print(f"⚠ Warning creating index {idx_name}: {e}")

        # Verify final state
        result = conn.execute(text("""
            SELECT column_name
            FROM information_schema.columns
            WHERE table_name='serial_numbers'
            ORDER BY column_name
        """))
        all_cols = [row[0] for row in result]

        print("\n✓ Migration complete!")
        print(f"✓ serial_numbers table now has {len(all_cols)} columns")
        print(f"✓ New columns: batch_number, branch_id")


if __name__ == "__main__":
    main()
