"""
MeshBay Node keystore — encrypted local storage for node identity keys.

Stores Ed25519 + X25519 keypairs and the GEK (Group Encryption Key)
encrypted at rest with AES-256-GCM, key derived via Argon2id.

Three unlock modes (checked in order):
  1. MESHBAY_UNLOCK_KEY env var  — for headless/systemd deployments
  2. unlock.key file             — lazy mode, chmod 600, documented risk
  3. Interactive prompt          — default secure mode

Keystore file format (JSON):
  {
    "version": 1,
    "argon2_salt_b64": "...",
    "iv_b64": "...",
    "tag_b64": "...",
    "ciphertext_b64": "..."
  }

Plaintext payload (msgpack, inside the AES-256-GCM envelope):
  {
    "sk_ed25519_b64": "...",
    "sk_x25519_b64":  "...",
    "gek_b64":        "..."   # may be absent until group is joined
  }
"""

import base64
import getpass
import json
import logging
import os
import sys
from dataclasses import dataclass
from pathlib import Path

import msgpack
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey

from meshbay_node.platform import chmod_private, config_dir
from meshbay_common.crypto import (
    ARGON2_ITERATIONS,
    ARGON2_LANES,
    ARGON2_MEMORY_COST,
    LEGACY_ARGON2_ITERATIONS,
    LEGACY_ARGON2_LANES,
    LEGACY_ARGON2_MEMORY_COST,
    decrypt_keystore,
    derive_keystore_key,
    encrypt_keystore,
    generate_gek,
    pk_to_b64,
    sk_to_b64,
    sk_to_raw,
)

log = logging.getLogger(__name__)

KEYSTORE_VERSION = 1
DEFAULT_KEYSTORE_PATH = config_dir() / "keystore.enc"
DEFAULT_UNLOCK_FILE   = config_dir() / "unlock.key"


@dataclass
class NodeKeys:
    sk_ed25519: Ed25519PrivateKey
    sk_x25519:  X25519PrivateKey
    gek:        bytes | None = None   # None until group is joined

    @property
    def pk_ed25519_b64(self) -> str:
        return pk_to_b64(self.sk_ed25519.public_key())

    @property
    def pk_x25519_b64(self) -> str:
        return pk_to_b64(self.sk_x25519.public_key())


# ── Password resolution ───────────────────────────────────────────────────────

def _resolve_password(unlock_file: Path | None = None) -> str:
    """
    Resolve the keystore password from (in order):
      1. MESHBAY_UNLOCK_KEY environment variable
      2. unlock.key file (if it exists and chmod 600)
      3. Interactive getpass prompt
    """
    # 1. Environment variable (systemd EnvironmentFile= pattern)
    env_key = os.environ.get("MESHBAY_UNLOCK_KEY")
    if env_key:
        log.debug("Keystore password from MESHBAY_UNLOCK_KEY")
        return env_key

    # 2. Unlock key file
    key_file = unlock_file or DEFAULT_UNLOCK_FILE
    if key_file.exists():
        if sys.platform != "win32":
            mode = oct(key_file.stat().st_mode)[-3:]
            if mode != "600":
                log.warning(
                    "unlock.key permissions are %s (expected 600) — fix with: "
                    "chmod 600 %s",
                    mode, key_file,
                )
        log.debug("Keystore password from %s", key_file)
        return key_file.read_text().strip()

    # 3. Interactive prompt
    return getpass.getpass("MeshBay node keystore password: ")


# ── Keystore I/O ──────────────────────────────────────────────────────────────

def _serialize_keys(keys: NodeKeys) -> bytes:
    payload = {
        "sk_ed25519_b64": sk_to_b64(keys.sk_ed25519),
        "sk_x25519_b64":  sk_to_b64(keys.sk_x25519),
    }
    if keys.gek is not None:
        payload["gek_b64"] = base64.b64encode(keys.gek).decode()
    return msgpack.packb(payload, use_bin_type=True)

def _deserialize_keys(data: bytes) -> NodeKeys:
    payload = msgpack.unpackb(data, raw=False)
    sk_ed = Ed25519PrivateKey.from_private_bytes(
        base64.b64decode(payload["sk_ed25519_b64"]))
    sk_x = X25519PrivateKey.from_private_bytes(
        base64.b64decode(payload["sk_x25519_b64"]))
    gek = base64.b64decode(payload["gek_b64"]) if "gek_b64" in payload else None
    return NodeKeys(sk_ed25519=sk_ed, sk_x25519=sk_x, gek=gek)


def create_keystore(
    path: Path | None = None,
    password: str | None = None,
    unlock_file: Path | None = None,
) -> NodeKeys:
    """
    Generate a new node identity, encrypt it, and save to disk.
    Fails if the keystore already exists.
    """
    path = path or DEFAULT_KEYSTORE_PATH
    if path.exists():
        raise FileExistsError(f"Keystore already exists: {path} — use load_keystore()")

    path.parent.mkdir(parents=True, exist_ok=True)

    pwd = password or _resolve_password(unlock_file)
    if len(pwd) < 8:
        raise ValueError("Password must be at least 8 characters")

    keys = NodeKeys(
        sk_ed25519=Ed25519PrivateKey.generate(),
        sk_x25519=X25519PrivateKey.generate(),
        gek=None,
    )

    _write_keystore(path, keys, pwd)
    log.info("New keystore created at %s", path)
    return keys


def load_keystore(
    path: Path | None = None,
    password: str | None = None,
    unlock_file: Path | None = None,
) -> NodeKeys:
    """Load and decrypt an existing keystore."""
    path = path or DEFAULT_KEYSTORE_PATH
    if not path.exists():
        raise FileNotFoundError(f"Keystore not found: {path} — run: meshbay-node init")

    pwd = password or _resolve_password(unlock_file)
    envelope = json.loads(path.read_text())

    if envelope.get("version") != KEYSTORE_VERSION:
        raise ValueError(f"Unsupported keystore version: {envelope.get('version')}")

    salt = base64.b64decode(envelope["argon2_salt_b64"])
    iv   = base64.b64decode(envelope["iv_b64"])
    tag  = base64.b64decode(envelope["tag_b64"])
    ct   = base64.b64decode(envelope["ciphertext_b64"])

    # Envelopes written before M2 carry no parameters and used the 64 MB profile.
    params = envelope.get("argon2", {
        "iterations":  LEGACY_ARGON2_ITERATIONS,
        "memory_cost": LEGACY_ARGON2_MEMORY_COST,
        "lanes":       LEGACY_ARGON2_LANES,
    })
    aes_key = derive_keystore_key(
        pwd, salt,
        iterations=params.get("iterations"),
        memory_cost=params.get("memory_cost"),
        lanes=params.get("lanes"),
    )
    try:
        plaintext = decrypt_keystore(iv, ct, tag, aes_key)
    except Exception:
        raise ValueError("Wrong password or corrupted keystore")

    log.info("Keystore loaded from %s", path)
    return _deserialize_keys(plaintext)


def save_keystore(
    keys: NodeKeys,
    path: Path | None = None,
    password: str | None = None,
    unlock_file: Path | None = None,
) -> None:
    """Re-encrypt and save updated keys (e.g. after GEK is set)."""
    path = path or DEFAULT_KEYSTORE_PATH
    pwd = password or _resolve_password(unlock_file)
    _write_keystore(path, keys, pwd)
    log.debug("Keystore updated at %s", path)


def _write_keystore(path: Path, keys: NodeKeys, password: str) -> None:
    salt = os.urandom(16)
    aes_key = derive_keystore_key(password, salt)
    plaintext = _serialize_keys(keys)
    iv, ct, tag = encrypt_keystore(plaintext, aes_key)

    envelope = {
        "version":         KEYSTORE_VERSION,
        "argon2_salt_b64": base64.b64encode(salt).decode(),
        # Recorded so parameters can be raised later without orphaning this file.
        "argon2": {
            "iterations":  ARGON2_ITERATIONS,
            "memory_cost": ARGON2_MEMORY_COST,
            "lanes":       ARGON2_LANES,
        },
        "iv_b64":          base64.b64encode(iv).decode(),
        "tag_b64":         base64.b64encode(tag).decode(),
        "ciphertext_b64":  base64.b64encode(ct).decode(),
    }
    path.write_text(json.dumps(envelope, indent=2))
    chmod_private(path)


def load_or_create_keystore(
    path: Path | None = None,
    password: str | None = None,
    unlock_file: Path | None = None,
) -> NodeKeys:
    """Load if exists, create if not. Convenience for daemon startup."""
    path = path or DEFAULT_KEYSTORE_PATH
    if path.exists():
        return load_keystore(path, password, unlock_file)
    return create_keystore(path, password, unlock_file)