"""
Brynq Credential Vault
Securely stores AI platform API keys using machine-specific encryption.

Encryption: Fernet (AES-128-CBC + HMAC-SHA256) via the `cryptography` library
when available. Falls back to PBKDF2 + HMAC-SHA256-CTR stream cipher (stdlib
only) if `cryptography` is not installed.

Key derivation: PBKDF2-HMAC-SHA256 (100K iterations) from a machine-specific
fingerprint (hostname + username). Keys are tied to the machine they were
created on.

Storage:
  ~/.brynq/vault.json — encrypted API keys, one per platform
"""

import base64
import getpass
import hashlib
import hmac
import json
import os
import secrets
import socket
from pathlib import Path
from typing import Optional

VAULT_FILE = Path.home() / ".brynq" / "vault.json"

# Key derivation parameters
_KDF_ITERATIONS = 100_000
_SALT_LEN = 16
_KEY_LEN = 32  # 256-bit derived key

# Version tags embedded in the encrypted payload header
_V3_TAG = b"BV3:"  # Fernet-based (v3)
_V2_HEADER_LEN = _SALT_LEN + 16 + 32  # salt + iv + mac (legacy v2)

# Try to import Fernet for proper AES encryption
try:
    from cryptography.fernet import Fernet, InvalidToken
    _HAS_FERNET = True
except ImportError:
    _HAS_FERNET = False


def _get_machine_fingerprint() -> bytes:
    """Derive a stable master secret from machine characteristics."""
    hostname = socket.gethostname()
    username = getpass.getuser()
    raw = f"{hostname}:{username}:brynq-vault-v3"
    return hashlib.sha256(raw.encode("utf-8")).digest()


def _derive_key(salt: bytes) -> bytes:
    """Derive an encryption key using PBKDF2-HMAC-SHA256."""
    master = _get_machine_fingerprint()
    return hashlib.pbkdf2_hmac(
        "sha256", master, salt, iterations=_KDF_ITERATIONS, dklen=_KEY_LEN
    )


def _derive_fernet_key(salt: bytes) -> bytes:
    """Derive a 32-byte key and encode it as url-safe base64 for Fernet.

    Fernet requires a 32-byte key that is url-safe base64-encoded.
    """
    raw_key = _derive_key(salt)
    return base64.urlsafe_b64encode(raw_key)


# ── Fernet encryption (v3) ───────────────────────────────────────────────


def _encrypt_fernet(plaintext: str) -> str:
    """Encrypt using Fernet (AES-128-CBC + HMAC-SHA256).

    Format: base64(_V3_TAG + salt + fernet_token)
    """
    salt = secrets.token_bytes(_SALT_LEN)
    fernet_key = _derive_fernet_key(salt)
    f = Fernet(fernet_key)
    token = f.encrypt(plaintext.encode("utf-8"))
    packed = _V3_TAG + salt + token
    return base64.b64encode(packed).decode("utf-8")


def _decrypt_fernet(packed: bytes) -> Optional[str]:
    """Decrypt a v3 (Fernet) vault entry. Returns None on failure."""
    if not packed.startswith(_V3_TAG):
        return None
    payload = packed[len(_V3_TAG):]
    if len(payload) < _SALT_LEN + 1:
        return None
    salt = payload[:_SALT_LEN]
    token = payload[_SALT_LEN:]
    fernet_key = _derive_fernet_key(salt)
    f = Fernet(fernet_key)
    try:
        return f.decrypt(token).decode("utf-8")
    except (InvalidToken, Exception):
        return None


# ── HMAC-CTR stream cipher fallback (v2, stdlib only) ────────────────────


def _hmac_keystream(key: bytes, iv: bytes, length: int) -> bytes:
    """Generate a pseudorandom keystream using HMAC-SHA256 in counter mode."""
    stream = bytearray()
    counter = 0
    while len(stream) < length:
        block_input = iv + counter.to_bytes(4, "big")
        block = hmac.new(key, block_input, hashlib.sha256).digest()
        stream.extend(block)
        counter += 1
    return bytes(stream[:length])


def _encrypt_hmac_ctr(plaintext: str) -> str:
    """Encrypt using HMAC-SHA256-CTR stream cipher with encrypt-then-MAC.

    Format: base64(salt + iv + mac + ciphertext)
    Used as fallback when `cryptography` is not installed.
    """
    salt = secrets.token_bytes(_SALT_LEN)
    iv = secrets.token_bytes(16)
    key = _derive_key(salt)

    pt_bytes = plaintext.encode("utf-8")
    keystream = _hmac_keystream(key, iv, len(pt_bytes))
    ciphertext = bytes(a ^ b for a, b in zip(pt_bytes, keystream))

    mac = hmac.new(key, ciphertext, hashlib.sha256).digest()
    packed = salt + iv + mac + ciphertext
    return base64.b64encode(packed).decode("utf-8")


def _decrypt_hmac_ctr(packed: bytes) -> Optional[str]:
    """Decrypt a v2 (HMAC-CTR) vault entry. Returns None on failure."""
    if len(packed) < _V2_HEADER_LEN:
        return None

    salt = packed[:_SALT_LEN]
    iv = packed[_SALT_LEN:_SALT_LEN + 16]
    mac_stored = packed[_SALT_LEN + 16:_V2_HEADER_LEN]
    ct = packed[_V2_HEADER_LEN:]

    key = _derive_key(salt)

    mac_computed = hmac.new(key, ct, hashlib.sha256).digest()
    if not hmac.compare_digest(mac_stored, mac_computed):
        return None

    keystream = _hmac_keystream(key, iv, len(ct))
    pt_bytes = bytes(a ^ b for a, b in zip(ct, keystream))

    try:
        return pt_bytes.decode("utf-8")
    except UnicodeDecodeError:
        return None


# ── Legacy XOR decryption (v1, read-only for migration) ─────────────────


def _decrypt_legacy_xor(ciphertext_b64: str) -> Optional[str]:
    """Fallback: decrypt entries encrypted with the original XOR-only scheme."""
    try:
        hostname = socket.gethostname()
        username = getpass.getuser()
        key = hashlib.sha256(f"{hostname}:{username}:brynq-salt".encode("utf-8")).digest()
        ct_bytes = base64.b64decode(ciphertext_b64.encode("utf-8"))
        decrypted = bytearray()
        for i, b in enumerate(ct_bytes):
            decrypted.append(b ^ key[i % len(key)])
        return decrypted.decode("utf-8")
    except Exception:
        return None


# ── Unified encrypt/decrypt dispatch ─────────────────────────────────────


def _encrypt(plaintext: str) -> str:
    """Encrypt a string using the best available method.

    Prefers Fernet (v3) when `cryptography` is installed, otherwise
    falls back to HMAC-CTR (v2).
    """
    if _HAS_FERNET:
        return _encrypt_fernet(plaintext)
    return _encrypt_hmac_ctr(plaintext)


def _decrypt(ciphertext_b64: str) -> str:
    """Decrypt a vault entry, auto-detecting the encryption version.

    Tries v3 (Fernet) -> v2 (HMAC-CTR) -> v1 (legacy XOR).
    Returns empty string on total failure.
    """
    try:
        packed = base64.b64decode(ciphertext_b64.encode("utf-8"))
    except Exception:
        # Not valid base64 — try legacy as last resort
        return _decrypt_legacy_xor(ciphertext_b64) or ""

    # Try v3 (Fernet) first
    if _HAS_FERNET and packed.startswith(_V3_TAG):
        result = _decrypt_fernet(packed)
        if result is not None:
            return result

    # Try v2 (HMAC-CTR)
    result = _decrypt_hmac_ctr(packed)
    if result is not None:
        return result

    # Try v1 (legacy XOR)
    return _decrypt_legacy_xor(ciphertext_b64) or ""


def _needs_migration(ciphertext_b64: str) -> bool:
    """Check if a stored value should be re-encrypted with the current scheme."""
    try:
        packed = base64.b64decode(ciphertext_b64.encode("utf-8"))
    except Exception:
        return True  # Can't even decode — definitely legacy

    # If Fernet is available and it's already v3, no migration needed
    if _HAS_FERNET and packed.startswith(_V3_TAG):
        return False

    # If Fernet is available but entry is v2 or v1 — migrate
    if _HAS_FERNET:
        return True

    # No Fernet — check if it's already v2 (HMAC-CTR)
    if len(packed) >= _V2_HEADER_LEN:
        salt = packed[:_SALT_LEN]
        iv = packed[_SALT_LEN:_SALT_LEN + 16]
        mac_stored = packed[_SALT_LEN + 16:_V2_HEADER_LEN]
        ct = packed[_V2_HEADER_LEN:]
        key = _derive_key(salt)
        mac_computed = hmac.new(key, ct, hashlib.sha256).digest()
        if hmac.compare_digest(mac_stored, mac_computed):
            return False  # Already v2, and no Fernet to upgrade to

    return True  # Legacy v1


# ── Vault I/O ────────────────────────────────────────────────────────────


def _load_vault() -> dict:
    """Load the encrypted vault from disk."""
    if not VAULT_FILE.exists():
        return {}
    try:
        data = json.loads(VAULT_FILE.read_text("utf-8"))
        return data.get("keys", {})
    except (json.JSONDecodeError, OSError):
        return {}


def _save_vault(keys: dict) -> None:
    """Save the encrypted vault to disk."""
    VAULT_FILE.parent.mkdir(parents=True, exist_ok=True)
    VAULT_FILE.write_text(json.dumps({"keys": keys}, indent=2) + "\n", encoding="utf-8")
    try:
        if os.name != "nt":
            VAULT_FILE.chmod(0o600)
    except Exception:
        pass


# ── Public API ───────────────────────────────────────────────────────────


def vault_store(platform: str, key: str) -> None:
    """Store an API key securely."""
    keys = _load_vault()
    keys[platform.lower()] = _encrypt(key)
    _save_vault(keys)


def vault_get(platform: str) -> Optional[str]:
    """Retrieve an API key. Auto-migrates legacy entries to current encryption."""
    keys = _load_vault()
    enc_key = keys.get(platform.lower())
    if not enc_key:
        return None
    result = _decrypt(enc_key)
    if result and _needs_migration(enc_key):
        # Re-encrypt with current best scheme
        keys[platform.lower()] = _encrypt(result)
        _save_vault(keys)
    return result or None


def vault_list() -> list:
    """List all connected platforms."""
    return list(_load_vault().keys())


def vault_delete(platform: str) -> bool:
    """Delete an API key."""
    keys = _load_vault()
    if platform.lower() in keys:
        del keys[platform.lower()]
        _save_vault(keys)
        return True
    return False
