"""Ed25519 keypair management for Muse CLI authentication.

Private keys are stored in ``~/.muse/keys/``, one PEM file per hub hostname
(or per ``hostname__agent_id`` for agent-specific keys).
The private key is PKCS8-encoded PEM (no passphrase) with strict 0o600
permissions, written atomically to prevent partial-write corruption.

Agent key paths
---------------
Human:  ``~/.muse/keys/localhost_10003.pem``
Agent:  ``~/.muse/keys/localhost_10003__agentception-abc123.pem``

The double-underscore ``__`` separates hostname from agent_id in the filename.
Forward-slashes and colons in the hostname are replaced with ``_`` as usual;
the agent_id portion is appended verbatim (it is a safe handle string).

Public key encoding
--------------------
Public keys are transmitted as URL-safe base64 (no padding) of the raw 32-byte
Ed25519 public key material — identical to the format expected by MuseHub's
``/api/auth/challenge`` and ``/api/auth/verify`` endpoints.

Fingerprint
-----------
SHA-256 hex digest of the raw 32 public key bytes, lowercase, no separators.
This is the stable identifier used by MuseHub to look up a registered key.

Security properties
-------------------
- Private keys are never logged, printed, or included in exception messages.
- Key files are written with ``fchmod(0o600)`` *before* any data is written,
  eliminating the TOCTOU window that ``write_text()`` + ``chmod()`` creates.
- Writes are atomic: data goes to a temp file in the same directory, then
  ``os.replace()`` renames it over the target.
- Symlink guard: if the target path is already a symlink, write is refused.
- The ``~/.muse/keys/`` directory is created with mode 0o700 (user-only).
"""
from __future__ import annotations

import base64
import hashlib
import logging
import os
import pathlib
import stat
import tempfile

from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives.serialization import (
    Encoding,
    NoEncryption,
    PrivateFormat,
    PublicFormat,
)

logger = logging.getLogger(__name__)

_KEYS_DIR = pathlib.Path.home() / ".muse" / "keys"


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _hostname_key(hostname: str) -> str:
    """Sanitise *hostname* to a safe filename component.

    Replaces ``:`` and ``/`` with ``_`` so that ``localhost:10003`` becomes a
    valid filename without path-traversal risk.
    """
    return hostname.replace(":", "_").replace("/", "_").lower()


def _key_path(hostname: str, agent_id: str | None = None) -> pathlib.Path:
    """Return the path to the PKCS8 PEM private key file for *hostname*.

    Agent keys use a ``{hostname}__{agent_id}.pem`` filename so that human
    and agent keys never collide in ``~/.muse/keys/``.
    """
    base = _hostname_key(hostname)
    if agent_id:
        return _KEYS_DIR / f"{base}__{agent_id}.pem"
    return _KEYS_DIR / f"{base}.pem"


def _ensure_keys_dir() -> None:
    """Create ``~/.muse/keys/`` with 0o700 permissions if it does not exist."""
    _KEYS_DIR.mkdir(parents=True, exist_ok=True)
    try:
        os.chmod(_KEYS_DIR, stat.S_IRWXU)  # 0o700
    except OSError as exc:
        logger.warning("⚠️ Could not set permissions on %s: %s", _KEYS_DIR, exc)


def _write_private_key_pem(path: pathlib.Path, private_key: Ed25519PrivateKey) -> None:
    """Write *private_key* to *path* as an unencrypted PKCS8 PEM file.

    Security guarantees:
    - Refuses to write if *path* is a symlink.
    - Sets 0o600 via ``fchmod`` *before* any data is written.
    - Uses ``os.replace()`` for an atomic rename from a temp file.
    """
    if path.is_symlink():
        raise OSError(
            f"Security: {path} is a symlink. "
            "Refusing to write a private key to a symlink target."
        )

    pem_bytes: bytes = private_key.private_bytes(
        encoding=Encoding.PEM,
        format=PrivateFormat.PKCS8,
        encryption_algorithm=NoEncryption(),
    )

    fd, tmp_path_str = tempfile.mkstemp(dir=path.parent, prefix=".key-tmp-")
    tmp_path = pathlib.Path(tmp_path_str)
    try:
        os.fchmod(fd, stat.S_IRUSR | stat.S_IWUSR)  # 0o600 before any data
        with os.fdopen(fd, "wb") as fh:
            fh.write(pem_bytes)
        os.replace(tmp_path, path)
    except Exception:
        try:
            tmp_path.unlink(missing_ok=True)
        except OSError:
            pass
        raise


# ---------------------------------------------------------------------------
# Public key encoding helpers
# ---------------------------------------------------------------------------


def public_key_to_b64url(public_key: Ed25519PublicKey) -> str:
    """Return the URL-safe base64 (no padding) encoding of the raw public key bytes."""
    raw: bytes = public_key.public_bytes(Encoding.Raw, PublicFormat.Raw)
    return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii")


def public_key_fingerprint(public_key: Ed25519PublicKey) -> str:
    """Return the SHA-256 hex fingerprint (64 lowercase hex chars) of the raw public key."""
    raw: bytes = public_key.public_bytes(Encoding.Raw, PublicFormat.Raw)
    return hashlib.sha256(raw).hexdigest()


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def generate_keypair(hostname: str, agent_id: str | None = None) -> tuple[str, str]:
    """Generate a new Ed25519 keypair and persist the private key to disk.

    Human keys are written to ``~/.muse/keys/{hostname}.pem`` (0o600).
    Agent keys are written to ``~/.muse/keys/{hostname}__{agent_id}.pem``.

    If a key already exists at the target path it is overwritten — callers
    should warn the user before invoking this function if the key exists.

    Args:
        hostname: Normalised hub hostname (e.g. ``"localhost:10003"``).
        agent_id: Optional agent handle.  When supplied the key is stored
                  under the agent-specific filename.

    Returns:
        ``(public_key_b64url, fingerprint)`` — both are safe to log / display.
        The private key is NEVER returned; it lives only on disk.
    """
    _ensure_keys_dir()
    private_key = Ed25519PrivateKey.generate()
    public_key = private_key.public_key()

    path = _key_path(hostname, agent_id)
    _write_private_key_pem(path, private_key)
    logger.info("✅ Ed25519 private key written to %s", path)

    return public_key_to_b64url(public_key), public_key_fingerprint(public_key)


def load_private_key(hostname: str, agent_id: str | None = None) -> Ed25519PrivateKey | None:
    """Load the Ed25519 private key for *hostname* (and optional *agent_id*) from disk.

    Args:
        hostname: Normalised hub hostname.
        agent_id: Optional agent handle — loads the agent-specific key file.

    Returns:
        :class:`Ed25519PrivateKey` if a key file exists, else ``None``.
        Returns ``None`` (not an exception) when no key has been generated
        yet — callers should prompt the user to run ``muse auth keygen``.
    """
    from cryptography.hazmat.primitives.serialization import load_pem_private_key

    path = _key_path(hostname, agent_id)
    if not path.is_file():
        return None
    try:
        pem_bytes = path.read_bytes()
        key = load_pem_private_key(pem_bytes, password=None)
    except Exception as exc:
        logger.warning("⚠️ Could not load private key from %s (%s)", path, type(exc).__name__)
        return None
    if not isinstance(key, Ed25519PrivateKey):
        logger.warning("⚠️ Key at %s is not an Ed25519 key", path)
        return None
    return key


def sign_bytes(private_key: Ed25519PrivateKey, data: bytes) -> str:
    """Sign *data* with *private_key* and return URL-safe base64 (no padding).

    Args:
        private_key: Ed25519 private key.
        data:        Raw bytes to sign (the nonce bytes from the challenge token).

    Returns:
        URL-safe base64 (no padding) of the 64-byte signature.
    """
    signature: bytes = private_key.sign(data)
    return base64.urlsafe_b64encode(signature).rstrip(b"=").decode("ascii")


def key_path_for(hostname: str, agent_id: str | None = None) -> pathlib.Path:
    """Return the expected private key path for *hostname* (may not exist).

    Args:
        hostname: Normalised hub hostname.
        agent_id: Optional agent handle — returns the agent-specific key path.
    """
    return _key_path(hostname, agent_id)


def load_private_key_from_pem(pem_bytes: bytes) -> Ed25519PrivateKey | None:
    """Load an Ed25519 private key from raw PEM bytes (e.g. from MUSE_AGENT_KEY env).

    Returns ``None`` if the bytes cannot be decoded or are not Ed25519.
    The bytes are never logged.
    """
    from cryptography.hazmat.primitives.serialization import load_pem_private_key

    try:
        key = load_pem_private_key(pem_bytes, password=None)
    except Exception as exc:
        logger.warning("⚠️ Could not load private key from PEM bytes (%s)", type(exc).__name__)
        return None
    if not isinstance(key, Ed25519PrivateKey):
        logger.warning("⚠️ PEM key is not Ed25519")
        return None
    return key