"""
Migrate communications table to use assetid instead of machineid.

This script updates the communications table FK from machineid to assetid.
Since assetid matches the original machineid, this is mostly a schema update.

Usage:
    python -m scripts.migration.migrate_communications
"""

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

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def run_migration(conn_str, dry_run=False):
    """
    Update communications to use assetid.

    Args:
        conn_str: Database connection string
        dry_run: If True, don't commit changes
    """
    engine = create_engine(conn_str)
    Session = sessionmaker(bind=engine)
    session = Session()

    try:
        # Check if assetid column already exists
        result = session.execute(text("""
            SELECT COUNT(*) FROM information_schema.columns
            WHERE table_name = 'communications' AND column_name = 'assetid'
        """))
        has_assetid = result.scalar() > 0

        if not has_assetid:
            logger.info("Adding assetid column to communications table...")

            # Add assetid column
            session.execute(text("""
                ALTER TABLE communications
                ADD COLUMN assetid INT NULL
            """))

            # Copy machineid values to assetid
            session.execute(text("""
                UPDATE communications
                SET assetid = machineid
                WHERE machineid IS NOT NULL
            """))

            # Add FK constraint (optional, depends on DB)
            try:
                session.execute(text("""
                    ALTER TABLE communications
                    ADD CONSTRAINT fk_comm_asset
                    FOREIGN KEY (assetid) REFERENCES assets(assetid)
                """))
            except Exception as e:
                logger.warning(f"Could not add FK constraint: {e}")

            logger.info("assetid column added and populated")
        else:
            logger.info("assetid column already exists")

        # Count records
        result = session.execute(text("""
            SELECT COUNT(*) FROM communications WHERE assetid IS NOT NULL
        """))
        count = result.scalar()
        logger.info(f"Communications with assetid: {count}")

        if dry_run:
            logger.info("Dry run - rolling back changes")
            session.rollback()
        else:
            session.commit()
            logger.info("Migration complete")

    except Exception as e:
        logger.error(f"Migration error: {e}")
        session.rollback()
        raise
    finally:
        session.close()


def main():
    import argparse
    import os
    import sys

    parser = argparse.ArgumentParser(description='Migrate communications to use assetid')
    parser.add_argument('--connection', help='Database connection string')
    parser.add_argument('--dry-run', action='store_true', help='Dry run without committing')

    args = parser.parse_args()

    conn_str = args.connection
    if not conn_str:
        sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
        from shopdb import create_app
        app = create_app()
        conn_str = app.config['SQLALCHEMY_DATABASE_URI']

    run_migration(conn_str, args.dry_run)


if __name__ == '__main__':
    main()