"""Migrate legacy plaintext memories to AES-256-GCM encryption.

Pre-0.5.0 builds stored memory content as raw UTF-8 bytes in the
``encrypted_content`` BLOB column with ``encryption_key_ref`` set to the
literal string ``mvp-plaintext``. This script walks every project DB and
replaces each plaintext row with a fresh per-memory AES-GCM ciphertext,
updating ``encryption_key_ref`` to the new key ID.

Migration is idempotent: rows that already have a real key reference
(anything other than ``mvp-plaintext``) are skipped.

Key management during migration:
    - One fresh master key is derived PER project (same logic as the
      runtime engine: PBKDF2 from machine_id + project_hash + salt).
    - Per-memory AES-256 keys are generated and persisted into the
      project's own ``encryption_keys`` table, encrypted with that
      project's master key.
    - No shared singleton — each project is a fully isolated context.

Usage:
    python -m soagila.core.encryption.migrate_plaintext [--dry-run]
    python -m soagila.core.encryption.migrate_plaintext --all-projects
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import os
import sys
from pathlib import Path
from uuid import uuid4

import aiosqlite
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes

logger = logging.getLogger("soagila.migrate_plaintext")

IV_SIZE = 12
KEY_SIZE = 32
SALT_SIZE = 32
PBKDF2_ITERATIONS = 480_000


def _derive_master_key(machine_id: str, project_hash: str, salt: bytes) -> bytes:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=KEY_SIZE,
        salt=salt,
        iterations=PBKDF2_ITERATIONS,
    )
    return kdf.derive(f"{machine_id}:{project_hash}".encode("utf-8"))


def _machine_id() -> str:
    from uuid import getnode
    user = os.environ.get("USERNAME", os.environ.get("USER", "soagila"))
    return f"{getnode()}-{user}"


def _load_or_create_salt(project_dir: Path) -> bytes:
    salt_path = project_dir / "master.salt"
    if salt_path.exists():
        return salt_path.read_bytes()
    salt = os.urandom(SALT_SIZE)
    salt_path.write_bytes(salt)
    return salt


def _aes_encrypt(plaintext: str, key: bytes) -> bytes:
    iv = os.urandom(IV_SIZE)
    aes = AESGCM(key)
    ct = aes.encrypt(iv, plaintext.encode("utf-8"), None)
    return iv + ct


def _wrap_key(raw_key: bytes, master_key: bytes) -> bytes:
    iv = os.urandom(IV_SIZE)
    aes = AESGCM(master_key)
    return iv + aes.encrypt(iv, raw_key, None)


async def _table_exists(db: aiosqlite.Connection, name: str) -> bool:
    async with db.execute(
        "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (name,)
    ) as cur:
        return (await cur.fetchone()) is not None


async def migrate_db(db_path: Path, dry_run: bool = False) -> dict:
    """Re-encrypt every plaintext row in a single project DB."""
    project_dir = db_path.parent
    project_hash = project_dir.name

    # Each project gets its own master key (no singleton!)
    salt = _load_or_create_salt(project_dir)
    master_key = _derive_master_key(_machine_id(), project_hash, salt)

    stats = {"episodic": 0, "semantic": 0, "skipped": 0, "no_table": 0}

    async with aiosqlite.connect(str(db_path)) as db:
        await db.execute("PRAGMA journal_mode=WAL")
        await db.execute("PRAGMA busy_timeout=30000")

        if not await _table_exists(db, "episodic_memories"):
            stats["no_table"] += 1
            return stats
        if not await _table_exists(db, "encryption_keys"):
            # Create the keys table so we can persist per-memory keys
            await db.execute(
                """
                CREATE TABLE IF NOT EXISTS encryption_keys (
                    key_ref TEXT PRIMARY KEY,
                    encrypted_key BLOB NOT NULL,
                    purpose TEXT,
                    created_at TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
                )
                """
            )

        # ── Episodic ────────────────────────────────────────────────
        async with db.execute(
            "SELECT id, encrypted_content, encryption_key_ref FROM episodic_memories"
        ) as cur:
            rows = await cur.fetchall()

        for mem_id, blob, key_ref in rows:
            if key_ref and key_ref != "mvp-plaintext":
                stats["skipped"] += 1
                continue
            try:
                plaintext = bytes(blob).decode("utf-8")
            except Exception:
                stats["skipped"] += 1
                continue

            raw_key = AESGCM.generate_key(bit_length=256)
            new_key_ref = f"key-memory-{uuid4().hex[:12]}"
            ciphertext = _aes_encrypt(plaintext, raw_key)
            wrapped = _wrap_key(raw_key, master_key)

            if dry_run:
                logger.info(
                    "[dry-run] episodic %s: %d -> %d bytes (key=%s)",
                    mem_id[:8], len(blob), len(ciphertext), new_key_ref,
                )
                continue

            await db.execute(
                "UPDATE episodic_memories SET encrypted_content = ?, encryption_key_ref = ? WHERE id = ?",
                (ciphertext, new_key_ref, mem_id),
            )
            await db.execute(
                "INSERT OR REPLACE INTO encryption_keys (key_ref, encrypted_key, purpose) VALUES (?, ?, ?)",
                (new_key_ref, wrapped, "memory"),
            )
            stats["episodic"] += 1

        # ── Semantic ────────────────────────────────────────────────
        if await _table_exists(db, "semantic_memories"):
            async with db.execute(
                "SELECT id, encrypted_content, encryption_key_ref FROM semantic_memories"
            ) as cur:
                rows = await cur.fetchall()

            for mem_id, blob, key_ref in rows:
                if key_ref and key_ref != "mvp-plaintext":
                    stats["skipped"] += 1
                    continue
                try:
                    plaintext = bytes(blob).decode("utf-8")
                except Exception:
                    stats["skipped"] += 1
                    continue

                raw_key = AESGCM.generate_key(bit_length=256)
                new_key_ref = f"key-memory-{uuid4().hex[:12]}"
                ciphertext = _aes_encrypt(plaintext, raw_key)
                wrapped = _wrap_key(raw_key, master_key)

                if dry_run:
                    logger.info(
                        "[dry-run] semantic %s: %d -> %d bytes (key=%s)",
                        mem_id[:8], len(blob), len(ciphertext), new_key_ref,
                    )
                    continue

                await db.execute(
                    "UPDATE semantic_memories SET encrypted_content = ?, encryption_key_ref = ? WHERE id = ?",
                    (ciphertext, new_key_ref, mem_id),
                )
                await db.execute(
                    "INSERT OR REPLACE INTO encryption_keys (key_ref, encrypted_key, purpose) VALUES (?, ?, ?)",
                    (new_key_ref, wrapped, "memory"),
                )
                stats["semantic"] += 1

        if not dry_run:
            await db.commit()

    return stats


async def main_async(args: argparse.Namespace) -> int:
    logging.basicConfig(
        level="DEBUG" if args.verbose else "INFO",
        format="[MIGRATE] %(asctime)s %(levelname)s %(message)s",
        stream=sys.stderr,
    )

    if args.all_projects:
        projects_dir = Path.home() / ".soagila" / "projects"
        db_paths = sorted(projects_dir.glob("*/soagila.db"))
    else:
        from soagila.config import get_project_db_path
        db_paths = [Path(get_project_db_path())]

    if not db_paths:
        logger.error("No project databases found")
        return 1

    total = {"episodic": 0, "semantic": 0, "skipped": 0, "no_table": 0}
    for db_path in db_paths:
        logger.info("Migrating %s", db_path)
        try:
            stats = await migrate_db(db_path, dry_run=args.dry_run)
        except Exception as exc:
            logger.error("  FAILED: %s", exc, exc_info=args.verbose)
            continue
        logger.info(
            "  episodic=%d semantic=%d skipped=%d no_table=%d",
            stats["episodic"], stats["semantic"], stats["skipped"], stats["no_table"],
        )
        for k, v in stats.items():
            total[k] += v

    logger.info("===== Migration complete =====")
    logger.info(
        "Total: episodic=%d semantic=%d skipped=%d no_table=%d (dry_run=%s)",
        total["episodic"], total["semantic"], total["skipped"],
        total["no_table"], args.dry_run,
    )
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--all-projects", action="store_true")
    parser.add_argument("-v", "--verbose", action="store_true")
    args = parser.parse_args()
    return asyncio.run(main_async(args))


if __name__ == "__main__":
    sys.exit(main())
