"""AES-256-GCM encryption layer for AgentForms.

Per-site encryption keys (AES-256-GCM). Master key from environment,
site-specific keys derived via HKDF from the master key and site ID.

Design:
- Each site gets its own encryption key derived from the master key.
- If the master key is not set, data is stored in plaintext (backwards compatible).
- If a site key is missing, it is lazily derived on first access.
- Encrypted values are stored as base64 in a single column: IV | tag | ciphertext.
- Key rotation: ENCRYPTION_KEY_V2 supersedes ENCRYPTION_KEY (V1).
  Decrypt tries V2 first, falls back to V1. Encrypt always uses V2.
"""

import base64
import hashlib
import os
import sqlite3

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes


# ─── Master key resolution (dual-key for rotation) ─────────────────────────────

def _resolve_master_keys():
    """Return (current_key, legacy_key) as bytes or None.

    Resolution:
    - ENCRYPTION_KEY_V2 set → current=V2, legacy=ENCRYPTION_KEY (if set and != V2)
    - ENCRYPTION_KEY_V2 not set → current=ENCRYPTION_KEY, legacy=None
    """
    v2 = os.environ.get("ENCRYPTION_KEY_V2")
    v1 = os.environ.get("ENCRYPTION_KEY")

    current_raw = v2 if v2 else v1
    legacy_raw = None
    if v2 and v1 and v1 != v2:
        legacy_raw = v1

    current = _raw_to_key(current_raw) if current_raw else None
    legacy = _raw_to_key(legacy_raw) if legacy_raw else None
    return current, legacy


def _raw_to_key(raw: str) -> bytes:
    """Convert raw env var string to 32-byte key."""
    if not raw:
        return None
    if all(c in "0123456789abcdefABCDEF" for c in raw) and len(raw) >= 64:
        return bytes.fromhex(raw[:64])
    return hashlib.sha256(raw.encode()).digest()


_master_keys_cache = {"current": None, "legacy": None}


def _get_master_key(version: str = "current") -> bytes:
    """Return the master encryption key. Raises RuntimeError if not configured."""
    if _master_keys_cache["current"] is None:
        _master_keys_cache["current"], _master_keys_cache["legacy"] = _resolve_master_keys()

    key = _master_keys_cache[version]
    if key is None:
        if version == "current":
            raise RuntimeError(
                "ENCRYPTION_KEY environment variable is required. "
                "Generate one with: python -c 'import secrets; print(secrets.token_hex(32))'"
            )
        return None  # legacy may legitimately be None
    return key


def _get_legacy_master_key() -> bytes | None:
    """Return the legacy (V1) master key, or None if not configured."""
    if _master_keys_cache["current"] is None:
        _master_keys_cache["current"], _master_keys_cache["legacy"] = _resolve_master_keys()
    return _master_keys_cache.get("legacy")


def master_key_is_rotation_active() -> bool:
    """Check if V2 key is configured (rotation mode)."""
    if _master_keys_cache["current"] is None:
        _master_keys_cache["current"], _master_keys_cache["legacy"] = _resolve_master_keys()
    return _master_keys_cache.get("legacy") is not None


# ─── Per-entity key derivation ──────────────────────────────────────────────────

def _derive_entity_key(master_key: bytes, scope: str, entity_id: int) -> bytes:
    """Derive a per-entity AES-256 key using HKDF (SHA-256)."""
    ikm = master_key
    salt = f"agentforms-{scope}-{entity_id}".encode()
    info = f"agentforms-aes256-{scope}".encode()
    kdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        info=info,
    )
    return kdf.derive(ikm)


# ─── Cache for site/user keys ──────────────────────────────────────────────────

_site_key_cache: dict[int, AESGCM] = {}
_user_key_cache: dict[int, AESGCM] = {}
_entity_key_cache: dict[str, AESGCM] = {}
_user_passkey_cache: dict[str, AESGCM] = {}


def _get_cipher_for_id(id_: int, scope: str) -> AESGCM:
    """Generic helper to derive a per-entity cipher (site, user, etc.) from master key."""
    cache = _site_key_cache if scope == "site" else _user_key_cache
    if id_ in cache:
        return cache[id_]

    master_key = _get_master_key("current")
    entity_key = _derive_entity_key(master_key, scope, id_)
    cipher = AESGCM(entity_key)
    cache[id_] = cipher
    return cipher


def _get_site_cipher(site_id: int) -> AESGCM:
    """Return AESGCM instance for the given site (current key)."""
    return _get_cipher_for_id(site_id, "site")


def _get_user_cipher(user_id: int) -> AESGCM:
    """Return AESGCM instance for the given user (current key)."""
    return _get_cipher_for_id(user_id, "user")


def _get_legacy_cipher_for_id(id_: int, scope: str) -> AESGCM | None:
    """Derive cipher using legacy master key, if available."""
    legacy_key = _get_legacy_master_key()
    if legacy_key is None:
        return None
    entity_key = _derive_entity_key(legacy_key, scope, id_)
    return AESGCM(entity_key)


def _get_legacy_site_cipher(site_id: int) -> AESGCM | None:
    """Return legacy AESGCM instance for the given site."""
    return _get_legacy_cipher_for_id(site_id, "site")


def _get_legacy_user_cipher(user_id: int) -> AESGCM | None:
    """Return legacy AESGCM instance for the given user."""
    return _get_legacy_cipher_for_id(user_id, "user")


# ─── Legacy site key derivation (info=b"agentforms-aes256") ────────────────────

def _derive_site_key(master_key: bytes, site_id: int) -> bytes:
    """Derive a per-site AES-256 key using HKDF (SHA-256) — legacy info string."""
    ikm = master_key
    salt = f"agentforms-site-{site_id}".encode()
    info = b"agentforms-aes256"
    kdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        info=info,
    )
    return kdf.derive(ikm)


def _get_legacy_site_cipher_old(site_id: int) -> AESGCM | None:
    """Return AESGCM using oldest derivation method (info=b'agentforms-aes256').

    This is the pre-refactor derivation. Used as a tertiary fallback.
    """
    try:
        master_key = _get_master_key("current")
    except RuntimeError:
        return None
    legacy_key = _derive_site_key(master_key, site_id)
    return AESGCM(legacy_key)


# ─── Password-based user key derivation ────────────────────────────────────────

def _derive_user_enc_key(password_hash: str) -> AESGCM | None:
    """Derive an AES-256-GCM cipher from a bcrypt password hash."""
    if not password_hash:
        return None

    if password_hash in _user_passkey_cache:
        return _user_passkey_cache[password_hash]

    ikm = password_hash.encode("utf-8")
    salt = b"agentforms-user-enc-v1"
    info = b"agentforms-submission-enc"
    kdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        info=info,
    )
    key = kdf.derive(ikm)
    cipher = AESGCM(key)
    _user_passkey_cache[password_hash] = cipher
    return cipher


# ─── Encrypt / Decrypt (site-scoped) ────────────────────────────────────────────

def encrypt_value(site_id: int, plaintext: str) -> str | None:
    """Encrypt a string value for the given site (current key)."""
    if not plaintext:
        return None

    cipher = _get_site_cipher(site_id)
    iv = os.urandom(12)
    ct = cipher.encrypt(iv, plaintext.encode("utf-8"), None)
    combined = iv + ct
    return base64.b64encode(combined).decode("ascii")


def decrypt_value(site_id: int, encrypted: str | None) -> str | None:
    """Decrypt a value for the given site.

    Tries in order:
    1. Current master key (V2 or V1 depending on config)
    2. Legacy master key (V1 if V2 is active)
    3. Oldest derivation method (info=b'agentforms-aes256')
    """
    if not encrypted:
        return None

    # Try 1: current key
    cipher = _get_site_cipher(site_id)
    result = _decrypt_with_cipher(cipher, encrypted)
    if result is not None:
        return result

    # Try 2: legacy master key
    legacy_cipher = _get_legacy_site_cipher(site_id)
    if legacy_cipher is not None:
        result = _decrypt_with_cipher(legacy_cipher, encrypted)
        if result is not None:
            return result

    # Try 3: oldest derivation (info=b'agentforms-aes256') with current key
    old_cipher = _get_legacy_site_cipher_old(site_id)
    if old_cipher is not None:
        result = _decrypt_with_cipher(old_cipher, encrypted)
        if result is not None:
            return result

    return None


# ─── Encrypt / Decrypt (user-scoped) ────────────────────────────────────────────

def encrypt_user_value(user_id: int, plaintext: str) -> str | None:
    """Encrypt a string value for the given user (user-scoped key, current)."""
    if not plaintext:
        return None

    cipher = _get_user_cipher(user_id)
    iv = os.urandom(12)
    ct = cipher.encrypt(iv, plaintext.encode("utf-8"), None)
    combined = iv + ct
    return base64.b64encode(combined).decode("ascii")


def decrypt_user_value(user_id: int, encrypted: str | None) -> str | None:
    """Decrypt a value for the given user (user-scoped key).

    Tries current key first, falls back to legacy key.
    """
    if not encrypted:
        return None

    cipher = _get_user_cipher(user_id)
    result = _decrypt_with_cipher(cipher, encrypted)
    if result is not None:
        return result

    legacy_cipher = _get_legacy_user_cipher(user_id)
    if legacy_cipher is not None:
        result = _decrypt_with_cipher(legacy_cipher, encrypted)
        if result is not None:
            return result

    return None


def try_decrypt_user_value(user_id: int, value: str | None) -> str | None:
    """Try to decrypt a user-scoped value, returning original if it fails."""
    if not value or not is_encrypted(value):
        return value
    result = decrypt_user_value(user_id, value)
    return result if result is not None else value


# ─── User-submission encryption (password-derived) ─────────────────────────────

def encrypt_user_submission(password_hash: str, plaintext: str) -> str | None:
    """Encrypt a submission field using the user's password-derived key."""
    if not plaintext:
        return None

    cipher = _derive_user_enc_key(password_hash)
    if not cipher:
        return None

    iv = os.urandom(12)
    ct = cipher.encrypt(iv, plaintext.encode("utf-8"), None)
    combined = iv + ct
    return base64.b64encode(combined).decode("ascii")


def decrypt_user_submission(password_hash: str, encrypted: str | None) -> str | None:
    """Decrypt a submission field using the user's password-derived key."""
    if not encrypted:
        return None

    cipher = _derive_user_enc_key(password_hash)
    if not cipher:
        return None

    return _decrypt_with_cipher(cipher, encrypted)


def try_decrypt_user_submission(password_hash: str, value: str | None, site_id: int | None = None) -> str | None:
    """Try to decrypt a submission value using user-scoped key first, then site-scoped fallback."""
    if not value or not is_encrypted(value):
        return value

    # First: try user-scoped decryption (new data)
    result = decrypt_user_submission(password_hash, value)
    if result is not None:
        return result

    # Fallback: try site-scoped decryption (old data)
    if site_id is not None:
        result = decrypt_value(site_id, value)
        if result is not None:
            return result

    return value


# ─── Entity-scoped encryption ──────────────────────────────────────────────────

def _get_entity_cipher(scope: str, entity_id: int) -> AESGCM:
    """Derive a per-entity cipher for non-site/non-user entities."""
    cache_key = f"{scope}-{entity_id}"
    if cache_key in _entity_key_cache:
        return _entity_key_cache[cache_key]

    master_key = _get_master_key("current")
    entity_key = _derive_entity_key(master_key, scope, entity_id)
    cipher = AESGCM(entity_key)
    _entity_key_cache[cache_key] = cipher
    return cipher


def _get_legacy_entity_cipher(scope: str, entity_id: int) -> AESGCM | None:
    """Derive entity cipher using legacy master key."""
    legacy_key = _get_legacy_master_key()
    if legacy_key is None:
        return None
    entity_key = _derive_entity_key(legacy_key, scope, entity_id)
    return AESGCM(entity_key)


def encrypt_entity(scope: str, entity_id: int, plaintext: str) -> str | None:
    """Encrypt a value using an entity-scoped key (current)."""
    if not plaintext:
        return None
    cipher = _get_entity_cipher(scope, entity_id)
    iv = os.urandom(12)
    ct = cipher.encrypt(iv, plaintext.encode("utf-8"), None)
    combined = iv + ct
    return base64.b64encode(combined).decode("ascii")


def decrypt_entity(scope: str, entity_id: int, encrypted: str | None) -> str | None:
    """Decrypt a value using an entity-scoped key.

    Tries current key first, falls back to legacy key.
    """
    if not encrypted:
        return None

    cipher = _get_entity_cipher(scope, entity_id)
    result = _decrypt_with_cipher(cipher, encrypted)
    if result is not None:
        return result

    legacy_cipher = _get_legacy_entity_cipher(scope, entity_id)
    if legacy_cipher is not None:
        result = _decrypt_with_cipher(legacy_cipher, encrypted)
        if result is not None:
            return result

    return None


def try_decrypt_entity(scope: str, entity_id: int, value: str | None) -> str | None:
    """Try to decrypt an entity-scoped value, returning original if it fails."""
    if not value or not is_encrypted(value):
        return value
    result = decrypt_entity(scope, entity_id, value)
    return result if result is not None else value


# ─── Helpers ─────────────────────────────────────────────────────────────────────

def _decrypt_with_cipher(cipher: AESGCM, encrypted: str) -> str | None:
    """Decrypt using a specific cipher instance."""
    try:
        combined = base64.b64decode(encrypted, validate=True)
        if len(combined) < 12 + 16:  # IV (12) + min ciphertext (1) + tag (16)
            return None
        iv = combined[:12]
        ct = combined[12:]
        plaintext = cipher.decrypt(iv, ct, None)
        return plaintext.decode("utf-8")
    except Exception:
        return None


def is_encrypted(value: str | None) -> bool:
    """Check if a value appears to be encrypted (base64-encoded with valid length)."""
    if not value:
        return False
    try:
        decoded = base64.b64decode(value, validate=True)
        return len(decoded) >= 12 + 1 + 16  # IV + min ciphertext + tag
    except Exception:
        return False


def try_decrypt(site_id: int, value: str | None) -> str | None:
    """Try to decrypt a value, returning the original if decryption fails."""
    if not value or not is_encrypted(value):
        return value
    result = decrypt_value(site_id, value)
    return result if result is not None else value


def hash_value(plaintext: str) -> str:
    """Compute a SHA-256 hash of a string for deterministic lookups."""
    if not plaintext:
        return ""
    return hashlib.sha256(plaintext.encode("utf-8")).hexdigest()


# ─── Key management ───────────────────────────────────────────────────────────

def key_is_configured() -> bool:
    """Check if encryption is configured (master key set)."""
    try:
        return _get_master_key("current") is not None
    except RuntimeError:
        return False


def clear_cache():
    """Clear all key caches. Call after rotation or env var changes."""
    _site_key_cache.clear()
    _user_key_cache.clear()
    _entity_key_cache.clear()
    _user_passkey_cache.clear()
    _master_keys_cache["current"] = None
    _master_keys_cache["legacy"] = None
    # Re-resolve
    _resolve_master_keys()


def reset_master_key_cache():
    """Force re-resolution of master keys (for testing/env changes)."""
    _master_keys_cache["current"] = None
    _master_keys_cache["legacy"] = None
