"""One-time migration that encrypts legacy plaintext OLT credentials."""
from sqlalchemy.orm.attributes import flag_modified

from app.core.database import SessionLocal
from app.models.device import OLTDevice
from app.core.credentials import CREDENTIAL_PREFIX


def migrate(batch_size: int = 100) -> int:
    """Encrypt every legacy credential and return the number of migrated records."""
    db = SessionLocal()
    migrated = 0
    try:
        devices = db.query(OLTDevice).yield_per(batch_size)
        for device in devices:
            if device.password.startswith(CREDENTIAL_PREFIX):
                continue
            # Reading a legacy row yields plaintext. Mark it dirty so the column type encrypts it on flush.
            device.password = device.password
            flag_modified(device, "password")
            migrated += 1
        db.commit()
        return migrated
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()


if __name__ == "__main__":
    print(f"Migrated {migrate()} OLT credential(s).")