"""
Migration script to add job reservations for the Fleet Calendar feature.

This creates two tables:
- job_reservations: Track future unit assignments for jobs/projects
- job_reservation_units: Link specific units to reservations

Run this script once to migrate an existing database.
"""

import sqlite3
import os

# Database path
DB_PATH = "./data/seismo_fleet.db"


def migrate_database():
    """Create the job_reservations and job_reservation_units tables"""

    if not os.path.exists(DB_PATH):
        print(f"Database not found at {DB_PATH}")
        print("The database will be created automatically when you run the application.")
        return

    print(f"Migrating database: {DB_PATH}")

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

    # Check if job_reservations table already exists
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='job_reservations'")
    if cursor.fetchone():
        print("Migration already applied - job_reservations table exists")
        conn.close()
        return

    print("Creating job_reservations table...")

    try:
        # Create job_reservations table
        cursor.execute("""
            CREATE TABLE job_reservations (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                project_id TEXT,
                start_date DATE NOT NULL,
                end_date DATE NOT NULL,
                assignment_type TEXT NOT NULL DEFAULT 'quantity',
                device_type TEXT DEFAULT 'seismograph',
                quantity_needed INTEGER,
                notes TEXT,
                color TEXT DEFAULT '#3B82F6',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        print("  Created job_reservations table")

        # Create indexes for job_reservations
        cursor.execute("CREATE INDEX idx_job_reservations_project_id ON job_reservations(project_id)")
        print("  Created index on project_id")

        cursor.execute("CREATE INDEX idx_job_reservations_dates ON job_reservations(start_date, end_date)")
        print("  Created index on dates")

        # Create job_reservation_units table
        print("Creating job_reservation_units table...")
        cursor.execute("""
            CREATE TABLE job_reservation_units (
                id TEXT PRIMARY KEY,
                reservation_id TEXT NOT NULL,
                unit_id TEXT NOT NULL,
                assignment_source TEXT DEFAULT 'specific',
                assigned_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (reservation_id) REFERENCES job_reservations(id),
                FOREIGN KEY (unit_id) REFERENCES roster(id)
            )
        """)
        print("  Created job_reservation_units table")

        # Create indexes for job_reservation_units
        cursor.execute("CREATE INDEX idx_job_reservation_units_reservation_id ON job_reservation_units(reservation_id)")
        print("  Created index on reservation_id")

        cursor.execute("CREATE INDEX idx_job_reservation_units_unit_id ON job_reservation_units(unit_id)")
        print("  Created index on unit_id")

        conn.commit()
        print("\nMigration completed successfully!")
        print("You can now use the Fleet Calendar to manage unit reservations.")

    except sqlite3.Error as e:
        print(f"\nError during migration: {e}")
        conn.rollback()
        raise

    finally:
        conn.close()


if __name__ == "__main__":
    migrate_database()
