#!/usr/bin/env python3
"""MAT-107: One-time migration to encrypt existing plaintext memory data.

Run INSIDE the memory-service container after deploying new code:
    docker exec -it matrix-ai-agent-memory-service-1 python migrate_encrypt.py

Connects as owner (memory) to bypass RLS.
"""
import os
import sys
import hashlib
import base64
import asyncio

import asyncpg
from cryptography.fernet import Fernet, InvalidToken

OWNER_DSN = os.environ.get(
    "OWNER_DATABASE_URL",
    "postgresql://memory:{password}@memory-db:5432/memories".format(
        password=os.environ.get("MEMORY_DB_OWNER_PASSWORD", "memory")
    ),
)
ENCRYPTION_KEY = os.environ.get("MEMORY_ENCRYPTION_KEY", "")


def _derive_user_key(user_id: str) -> bytes:
    derived = hashlib.pbkdf2_hmac("sha256", ENCRYPTION_KEY.encode(), user_id.encode(), iterations=1)
    return base64.urlsafe_b64encode(derived)


def _encrypt(text: str, user_id: str) -> str:
    f = Fernet(_derive_user_key(user_id))
    return f.encrypt(text.encode()).decode()


def _is_encrypted(text: str, user_id: str) -> bool:
    """Check if text is already Fernet-encrypted."""
    try:
        f = Fernet(_derive_user_key(user_id))
        f.decrypt(text.encode())
        return True
    except (InvalidToken, Exception):
        return False


async def migrate():
    if not ENCRYPTION_KEY:
        print("ERROR: MEMORY_ENCRYPTION_KEY not set")
        sys.exit(1)

    conn = await asyncpg.connect(OWNER_DSN)

    # Migrate memories
    rows = await conn.fetch("SELECT id, user_id, fact FROM memories ORDER BY id")
    print(f"Migrating {len(rows)} memories...")
    encrypted = 0
    skipped = 0
    for row in rows:
        if _is_encrypted(row["fact"], row["user_id"]):
            skipped += 1
            continue
        enc_fact = _encrypt(row["fact"], row["user_id"])
        await conn.execute("UPDATE memories SET fact = $1 WHERE id = $2", enc_fact, row["id"])
        encrypted += 1
        if encrypted % 100 == 0:
            print(f"  memories: {encrypted}/{len(rows)} encrypted")
    print(f"Memories done: {encrypted} encrypted, {skipped} already encrypted")

    # Migrate conversation_chunks
    rows = await conn.fetch("SELECT id, user_id, chunk_text, summary FROM conversation_chunks ORDER BY id")
    print(f"Migrating {len(rows)} chunks...")
    encrypted = 0
    skipped = 0
    for row in rows:
        if _is_encrypted(row["chunk_text"], row["user_id"]):
            skipped += 1
            continue
        enc_text = _encrypt(row["chunk_text"], row["user_id"])
        enc_summary = _encrypt(row["summary"], row["user_id"])
        await conn.execute(
            "UPDATE conversation_chunks SET chunk_text = $1, summary = $2 WHERE id = $3",
            enc_text, enc_summary, row["id"],
        )
        encrypted += 1
        if encrypted % 500 == 0:
            print(f"  chunks: {encrypted}/{len(rows)} encrypted")
    print(f"Chunks done: {encrypted} encrypted, {skipped} already encrypted")

    await conn.close()
    print("Migration complete.")


if __name__ == "__main__":
    asyncio.run(migrate())