"""
Migration: Add estimated_units to job_reservations

Adds column:
- job_reservations.estimated_units: Estimated number of units for the reservation (nullable integer)
"""

import sqlite3
import sys
from pathlib import Path

# Default database path (matches production pattern)
DB_PATH = "./data/seismo_fleet.db"


def migrate(db_path: str):
    """Run the migration."""
    print(f"Migrating database: {db_path}")

    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    try:
        # Check if job_reservations table exists
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='job_reservations'")
        if not cursor.fetchone():
            print("job_reservations table does not exist. Skipping migration.")
            return

        # Get existing columns in job_reservations
        cursor.execute("PRAGMA table_info(job_reservations)")
        existing_cols = {row[1] for row in cursor.fetchall()}

        # Add estimated_units column if it doesn't exist
        if 'estimated_units' not in existing_cols:
            print("Adding estimated_units column to job_reservations...")
            cursor.execute("ALTER TABLE job_reservations ADD COLUMN estimated_units INTEGER")
        else:
            print("estimated_units column already exists. Skipping.")

        conn.commit()
        print("Migration completed successfully!")

    except Exception as e:
        print(f"Migration failed: {e}")
        conn.rollback()
        raise
    finally:
        conn.close()


if __name__ == "__main__":
    db_path = DB_PATH

    if len(sys.argv) > 1:
        db_path = sys.argv[1]

    if not Path(db_path).exists():
        print(f"Database not found: {db_path}")
        sys.exit(1)

    migrate(db_path)