"""
Credential encryption and key rotation for AutoPIL catalog connections.

Encryption uses Fernet (symmetric AES-128-CBC + HMAC-SHA256).

Environment variables
---------------------
AUTOPIL_CATALOG_ENCRYPTION_KEY     Current encryption key (base64-urlsafe, 32 bytes).
AUTOPIL_CATALOG_ENCRYPTION_KEY_OLD Previous key accepted during a rotation window.
                                    Remove once all credentials have been re-encrypted.

Key rotation workflow
---------------------
1. Generate a new key:
       python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())"
2. Set both env vars:
       AUTOPIL_CATALOG_ENCRYPTION_KEY=<new_key>
       AUTOPIL_CATALOG_ENCRYPTION_KEY_OLD=<old_key>
3. Run rotation:
       autopil-admin catalog rotate-key
       (or call rotate_encryption_key() programmatically)
4. Once rotation completes successfully, remove AUTOPIL_CATALOG_ENCRYPTION_KEY_OLD.
"""

from __future__ import annotations

import json
import logging
import os
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from autopil.db.base import CatalogStoreBase

logger = logging.getLogger(__name__)


# ── key helpers ───────────────────────────────────────────────────────────────

def _current_key() -> str | None:
    return os.environ.get("AUTOPIL_CATALOG_ENCRYPTION_KEY")


def _old_key() -> str | None:
    return os.environ.get("AUTOPIL_CATALOG_ENCRYPTION_KEY_OLD")


def encrypt_credentials(data: dict) -> str:
    """
    Encrypt a credentials dict for storage.

    If AUTOPIL_CATALOG_ENCRYPTION_KEY is set, encrypts with Fernet.
    Otherwise stores as plaintext JSON (development / single-node mode).
    """
    raw = json.dumps(data)
    key = _current_key()
    if key:
        try:
            from cryptography.fernet import Fernet
            return Fernet(key.encode()).encrypt(raw.encode()).decode()
        except Exception as exc:
            raise ValueError(f"Failed to encrypt credentials: {exc}") from exc
    return raw


def decrypt_credentials(enc: str) -> dict:
    """
    Decrypt a stored credentials blob.

    Tries the current key first, then falls back to the old key (rotation window).
    Falls back to plaintext JSON if neither key is set.
    """
    current = _current_key()
    old = _old_key()

    if current:
        try:
            from cryptography.fernet import Fernet
            return json.loads(Fernet(current.encode()).decrypt(enc.encode()))
        except Exception:
            pass  # try old key

    if old:
        try:
            from cryptography.fernet import Fernet
            return json.loads(Fernet(old.encode()).decrypt(enc.encode()))
        except Exception:
            pass

    # No keys set, or both failed — assume plaintext JSON
    try:
        return json.loads(enc)
    except json.JSONDecodeError as exc:
        raise ValueError("Cannot decrypt credentials: no valid key and not plaintext JSON") from exc


# ── rotation ──────────────────────────────────────────────────────────────────

def rotate_encryption_key(
    store: "CatalogStoreBase",
    old_key: str | None = None,
    new_key: str | None = None,
    *,
    dry_run: bool = False,
) -> dict:
    """
    Re-encrypt all catalog connection credentials with a new Fernet key.

    Parameters
    ----------
    store:
        Any CatalogStoreBase implementation (SQLite or Postgres).
    old_key:
        Key currently used to encrypt stored credentials.
        Defaults to AUTOPIL_CATALOG_ENCRYPTION_KEY_OLD, then
        AUTOPIL_CATALOG_ENCRYPTION_KEY.
    new_key:
        Key to encrypt with after rotation.
        Defaults to AUTOPIL_CATALOG_ENCRYPTION_KEY.
    dry_run:
        If True, decrypt every credential and verify it can be re-encrypted,
        but do not write any changes.

    Returns
    -------
    dict with keys: rotated, skipped, failed, dry_run
    """
    resolved_old = old_key or _old_key() or _current_key()
    resolved_new = new_key or _current_key()

    if not resolved_new:
        raise ValueError(
            "No new key provided. Set AUTOPIL_CATALOG_ENCRYPTION_KEY or pass new_key."
        )
    if resolved_old == resolved_new:
        raise ValueError("old_key and new_key are identical — nothing to rotate.")

    try:
        from cryptography.fernet import Fernet
        old_fernet = Fernet(resolved_old.encode()) if resolved_old else None
        new_fernet = Fernet(resolved_new.encode())
    except Exception as exc:
        raise ValueError(f"Invalid Fernet key: {exc}") from exc

    # Collect all connections across all tenants (store.list_connections requires tenant_id,
    # so we use a low-level scan via the internal method if available, otherwise we use the
    # list_all_connections helper if present — fall back to per-known-tenant iteration).
    all_connections = _list_all_connections(store)

    rotated = 0
    skipped = 0
    failed = 0

    for conn in all_connections:
        try:
            # Decrypt with old key (or plaintext fallback)
            raw = conn.credentials_enc
            creds_dict = _decrypt_with_key(raw, old_fernet)

            # Re-encrypt with new key
            new_enc = new_fernet.encrypt(json.dumps(creds_dict).encode()).decode()

            if dry_run:
                # Just verify round-trip decryption
                _verify_round_trip(new_enc, new_fernet, creds_dict)
                logger.debug("dry_run: would rotate connection %s", conn.connection_id)
                rotated += 1
                continue

            # Write the re-encrypted blob back
            store.update_connection(
                conn.connection_id,
                {"credentials_enc": new_enc},
                conn.tenant_id,
            )
            rotated += 1
            logger.info("Rotated credentials for connection %s (tenant=%s)", conn.connection_id, conn.tenant_id)

        except _SkipConnection:
            skipped += 1
            logger.debug("Skipped connection %s (already plaintext or not encrypted)", conn.connection_id)
        except Exception as exc:
            failed += 1
            logger.error("Failed to rotate connection %s: %s", conn.connection_id, exc)
            if not dry_run:
                # Stop on first failure to avoid partial rotation
                raise RuntimeError(
                    f"Rotation aborted after {rotated} successes — "
                    f"failed on {conn.connection_id}: {exc}"
                ) from exc

    return {
        "rotated": rotated,
        "skipped": skipped,
        "failed": failed,
        "dry_run": dry_run,
        "total": len(all_connections),
    }


# ── internal helpers ──────────────────────────────────────────────────────────

class _SkipConnection(Exception):
    """Raised when a connection should be skipped (e.g. already plaintext)."""


def _decrypt_with_key(enc: str, fernet) -> dict:
    """Decrypt `enc` using `fernet`, or parse as plaintext JSON if no key."""
    if fernet:
        try:
            return json.loads(fernet.decrypt(enc.encode()))
        except Exception:
            # Try plaintext fallback (connections registered without a key)
            try:
                result = json.loads(enc)
                raise _SkipConnection("plaintext credential not encrypted")
            except _SkipConnection:
                raise
            except Exception:
                raise ValueError(f"Cannot decrypt credentials — wrong key or corrupt data")
    # No key provided — must be plaintext
    try:
        return json.loads(enc)
    except json.JSONDecodeError as exc:
        raise ValueError("No key provided and credential is not plaintext JSON") from exc


def _verify_round_trip(new_enc: str, new_fernet, original: dict) -> None:
    """Verify that new_enc decrypts back to original."""
    roundtripped = json.loads(new_fernet.decrypt(new_enc.encode()))
    if roundtripped != original:
        raise ValueError("Round-trip verification failed — decrypted value does not match original")


def _list_all_connections(store: "CatalogStoreBase"):
    """
    Best-effort scan of all connections regardless of tenant.

    SQLiteCatalogStore and PostgresCatalogStore both expose an internal
    _list_all_connections() if available; otherwise fall back to the
    public list_connections() which requires a tenant_id — in that case
    we can't iterate all tenants without a tenant registry, so we raise
    a helpful error.
    """
    if hasattr(store, "list_all_connections"):
        return store.list_all_connections()
    raise NotImplementedError(
        "This catalog store does not support listing all connections across tenants. "
        "Implement list_all_connections() on your CatalogStoreBase subclass, or "
        "pass a tenant-scoped store and call rotate_encryption_key() per tenant."
    )
