"""
Re-encrypt all stored HL API keys to the current enc:v2 format (H4 fix).

Upgrades, in place:
  * enc:v1 blobs (legacy unsalted-SHA256 KEK derivation)
  * plaintext rows (pre-encryption era)
to enc:v2 (PBKDF2-HMAC-SHA256, per-blob salt). Rows already in enc:v2 are
skipped — the script is idempotent and safe to re-run.

KEK rotation: set OLD_ENCRYPTION_KEY in the environment to decrypt with the
old KEK while encrypting with the current ENCRYPTION_KEY.

Usage (BACK UP THE DB FIRST):
    DATABASE_URL=<prod-url> python scripts/reencrypt_keys.py [--dry-run]
"""
import argparse
import asyncio
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy import select

from app.database import AsyncSessionLocal
from app.models import Subscription
from app.services import crypto


async def main(dry_run: bool) -> None:
    old_kek = os.environ.get("OLD_ENCRYPTION_KEY")
    if old_kek:
        # Decrypt with the old KEK, encrypt with the new one. Swapping the
        # setting around each call keeps crypto.py's caching simple.
        print("KEK rotation mode: decrypting with OLD_ENCRYPTION_KEY")

    upgraded = skipped = failed = 0
    async with AsyncSessionLocal() as db:
        rows = (await db.execute(
            select(Subscription).where(Subscription.hl_api_key.is_not(None))
        )).scalars().all()
        print(f"{len(rows)} subscription(s) with a stored HL key")

        for sub in rows:
            stored = sub.hl_api_key
            if crypto.is_current_format(stored) and not old_kek:
                skipped += 1
                continue
            try:
                if old_kek:
                    current = crypto.settings.encryption_key
                    crypto.settings.encryption_key = old_kek
                    crypto._fernet_v1 = None
                    crypto._fernet_v2_cache.clear()
                    try:
                        plaintext = crypto.decrypt_api_key(stored)
                    finally:
                        crypto.settings.encryption_key = current
                        crypto._fernet_v1 = None
                        crypto._fernet_v2_cache.clear()
                        crypto._encrypt_salt = None
                else:
                    plaintext = crypto.decrypt_api_key(stored)
                new_blob = crypto.encrypt_api_key(plaintext)
                # Round-trip check before touching the row.
                assert crypto.decrypt_api_key(new_blob) == plaintext
                if not dry_run:
                    sub.hl_api_key = new_blob
                upgraded += 1
                print(f"  {sub.wallet_address}: upgraded"
                      f"{' (dry-run)' if dry_run else ''}")
            except Exception as exc:
                failed += 1
                print(f"  {sub.wallet_address}: FAILED — {type(exc).__name__}: {exc}")

        if not dry_run and upgraded:
            await db.commit()
            print("Committed.")

    print(f"Done: upgraded={upgraded} skipped(already v2)={skipped} failed={failed}")
    if failed:
        sys.exit(1)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--dry-run", action="store_true",
                    help="report what would change without writing")
    args = ap.parse_args()
    asyncio.run(main(args.dry_run))