"""
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
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_common.crypto import (
    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 = Path.home() / ".config" / "meshbay" / "keystore.enc"
DEFAULT_UNLOCK_FILE   = Path.home() / ".config" / "meshbay" / "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():
        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"])

    aes_key = derive_keystore_key(pwd, salt)
    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(),
        "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))
    path.chmod(0o600)


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)