#!/usr/bin/env python3
"""
Script to rename a unit ID in the database.
This updates the unit across all tables with proper foreign key handling.
"""

import sys
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker

DATABASE_URL = "sqlite:///data/seismo_fleet.db"

def rename_unit(old_id: str, new_id: str):
    """
    Rename a unit ID across all relevant tables.

    Args:
        old_id: Current unit ID (e.g., "SLM4301")
        new_id: New unit ID (e.g., "SLM-43-01")
    """
    engine = create_engine(DATABASE_URL)
    Session = sessionmaker(bind=engine)
    session = Session()

    try:
        # Check if old unit exists
        result = session.execute(
            text("SELECT id, device_type FROM roster WHERE id = :old_id"),
            {"old_id": old_id}
        ).fetchone()

        if not result:
            print(f"❌ Error: Unit '{old_id}' not found in roster")
            return False

        device_type = result[1]
        print(f"✓ Found unit '{old_id}' (device_type: {device_type})")

        # Check if new ID already exists
        result = session.execute(
            text("SELECT id FROM roster WHERE id = :new_id"),
            {"new_id": new_id}
        ).fetchone()

        if result:
            print(f"❌ Error: Unit ID '{new_id}' already exists")
            return False

        print(f"\n🔄 Renaming '{old_id}' → '{new_id}'...\n")

        # Update roster table (primary)
        session.execute(
            text("UPDATE roster SET id = :new_id WHERE id = :old_id"),
            {"new_id": new_id, "old_id": old_id}
        )
        print(f"  ✓ Updated roster")

        # Update emitters table
        result = session.execute(
            text("UPDATE emitters SET id = :new_id WHERE id = :old_id"),
            {"new_id": new_id, "old_id": old_id}
        )
        if result.rowcount > 0:
            print(f"  ✓ Updated emitters ({result.rowcount} rows)")

        # Update unit_history table
        result = session.execute(
            text("UPDATE unit_history SET unit_id = :new_id WHERE unit_id = :old_id"),
            {"new_id": new_id, "old_id": old_id}
        )
        if result.rowcount > 0:
            print(f"  ✓ Updated unit_history ({result.rowcount} rows)")

        # Update deployed_with_modem_id references
        result = session.execute(
            text("UPDATE roster SET deployed_with_modem_id = :new_id WHERE deployed_with_modem_id = :old_id"),
            {"new_id": new_id, "old_id": old_id}
        )
        if result.rowcount > 0:
            print(f"  ✓ Updated modem references ({result.rowcount} rows)")

        # Update unit_assignments table (if exists)
        try:
            result = session.execute(
                text("UPDATE unit_assignments SET unit_id = :new_id WHERE unit_id = :old_id"),
                {"new_id": new_id, "old_id": old_id}
            )
            if result.rowcount > 0:
                print(f"  ✓ Updated unit_assignments ({result.rowcount} rows)")
        except Exception:
            pass  # Table may not exist

        # Update monitoring_sessions table (if exists)
        try:
            result = session.execute(
                text("UPDATE monitoring_sessions SET unit_id = :new_id WHERE unit_id = :old_id"),
                {"new_id": new_id, "old_id": old_id}
            )
            if result.rowcount > 0:
                print(f"  ✓ Updated monitoring_sessions ({result.rowcount} rows)")
        except Exception:
            pass  # Table may not exist

        # Commit all changes
        session.commit()
        print(f"\n✅ Successfully renamed unit '{old_id}' to '{new_id}'")
        return True

    except Exception as e:
        session.rollback()
        print(f"\n❌ Error during rename: {e}")
        return False
    finally:
        session.close()


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python rename_unit.py <old_id> <new_id>")
        print("Example: python rename_unit.py SLM4301 SLM-43-01")
        sys.exit(1)

    old_id = sys.argv[1]
    new_id = sys.argv[2]

    print(f"Unit Renaming Tool")
    print(f"=" * 50)
    print(f"Old ID: {old_id}")
    print(f"New ID: {new_id}")
    print(f"=" * 50)

    confirm = input(f"\nAre you sure you want to rename '{old_id}' to '{new_id}'? (yes/no): ")
    if confirm.lower() != 'yes':
        print("❌ Rename cancelled")
        sys.exit(0)

    success = rename_unit(old_id, new_id)
    sys.exit(0 if success else 1)
