"""
Migration: Add TBD date support to job reservations

Adds columns:
- job_reservations.estimated_end_date: For planning when end is TBD
- job_reservations.end_date_tbd: Boolean flag for TBD end dates
- job_reservation_units.unit_start_date: Unit-specific start (for swaps)
- job_reservation_units.unit_end_date: Unit-specific end (for swaps)
- job_reservation_units.unit_end_tbd: Unit-specific TBD flag
- job_reservation_units.notes: Notes for the assignment

Also makes job_reservations.end_date nullable.
"""

import sqlite3
import sys
from pathlib import Path

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 new columns to job_reservations if they don't exist
        if 'estimated_end_date' not in existing_cols:
            print("Adding estimated_end_date column to job_reservations...")
            cursor.execute("ALTER TABLE job_reservations ADD COLUMN estimated_end_date DATE")

        if 'end_date_tbd' not in existing_cols:
            print("Adding end_date_tbd column to job_reservations...")
            cursor.execute("ALTER TABLE job_reservations ADD COLUMN end_date_tbd BOOLEAN DEFAULT 0")

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

        # Add new columns to job_reservation_units if they don't exist
        if 'unit_start_date' not in unit_cols:
            print("Adding unit_start_date column to job_reservation_units...")
            cursor.execute("ALTER TABLE job_reservation_units ADD COLUMN unit_start_date DATE")

        if 'unit_end_date' not in unit_cols:
            print("Adding unit_end_date column to job_reservation_units...")
            cursor.execute("ALTER TABLE job_reservation_units ADD COLUMN unit_end_date DATE")

        if 'unit_end_tbd' not in unit_cols:
            print("Adding unit_end_tbd column to job_reservation_units...")
            cursor.execute("ALTER TABLE job_reservation_units ADD COLUMN unit_end_tbd BOOLEAN DEFAULT 0")

        if 'notes' not in unit_cols:
            print("Adding notes column to job_reservation_units...")
            cursor.execute("ALTER TABLE job_reservation_units ADD COLUMN notes TEXT")

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

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


if __name__ == "__main__":
    # Default to dev database
    db_path = "./data-dev/seismo_fleet.db"

    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)