"""Agent identity and commit signing for the Muse VCS.

Every commit in Muse can carry cryptographic provenance metadata that
identifies *who* or *what* produced it — a human author, an autonomous AI
agent, or a specific toolchain run.

Signing model
-------------
Signatures use **Ed25519** with the per-hub identity keypair stored under
``~/.muse/keys/{hostname}.pem`` (same keypair used for MSign request
authentication).  This is an asymmetric scheme: the private key signs; the
public key (embedded in the commit record) verifies.  Any party with the
commit record can verify the signature without access to the private key or
any external service.

Signature format
----------------
The signed input is :func:`provenance_payload` — a SHA-256 hex digest that
binds the commit content identity (``commit_id``) to authorship claims
(``author``, ``agent_id``, ``model_id``, ``toolchain_id``, ``prompt_hash``).

``CommitRecord.signature``
    Base64url-encoded Ed25519 signature (no padding), 86 characters.

``CommitRecord.signer_public_key``
    Base64url-encoded raw Ed25519 public key bytes (32 bytes → 43 chars).
    Embedded in the commit record so that verification is fully offline.

``CommitRecord.signer_key_id``
    First 16 hex characters of SHA-256(raw public key bytes).  Short enough
    to log, long enough for practical uniqueness.

Key management
--------------
Keys are the same Ed25519 keypairs used for MSign HTTP authentication.
Generate and register a keypair with::

    muse auth keygen   --hub http://localhost:10003
    muse auth register --hub http://localhost:10003 --handle <handle>

Usage
-----
::

    from muse.core.provenance import (
        make_agent_identity, sign_commit_ed25519, verify_commit_ed25519,
        public_key_fingerprint, sign_commit_record,
    )
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

    private_key: Ed25519PrivateKey = ...
    payload = provenance_payload(commit_id, agent_id="my-agent", ...)
    sig = sign_commit_ed25519(payload, private_key)
    pub_bytes = private_key.public_key().public_bytes(Raw, Raw)
    assert verify_commit_ed25519(payload, sig, pub_bytes)
"""

from __future__ import annotations

import base64
import hashlib
import logging
import pathlib

from typing import TYPE_CHECKING, TypedDict

if TYPE_CHECKING:
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

logger = logging.getLogger(__name__)

# Null-byte field separator — same convention as compute_commit_id in snapshot.py.
# Prevents separator-injection attacks where a field value contains the separator.
_PROV_SEP = "\x00"

# Version constant for the provenance payload format.
# v2 adds: version prefix + committed_at timestamp binding.
# v1 (no prefix): commit_id + author + agent_id + model_id + toolchain_id + prompt_hash
# v2 (this version): "muse-provenance-v2\n" + same fields + committed_at
PROVENANCE_PAYLOAD_VERSION = 2
_PROV_VERSION_PREFIX = "muse-provenance-v2\n"


# ---------------------------------------------------------------------------
# Provenance signing payload
# ---------------------------------------------------------------------------


def provenance_payload(
    commit_id: str,
    *,
    author: str = "",
    agent_id: str = "",
    model_id: str = "",
    toolchain_id: str = "",
    prompt_hash: str = "",
    committed_at: str = "",
) -> str:
    """Compute the SHA-256 hex digest of the provenance signing payload (v2).

    Binds the commit's *content identity* (``commit_id``, which covers
    snapshot + message + parents + timestamp) to its *authorship claims*
    (``author``, ``agent_id``, ``model_id``, ``toolchain_id``,
    ``prompt_hash``) and the ``committed_at`` wall-clock timestamp.

    **Version 2 changes (this version):**
    - Payload is prefixed with ``"muse-provenance-v2\\n"`` so v1 and v2
      signatures are clearly distinguishable.
    - ``committed_at`` is appended as the last field so that the commit
      timestamp cannot be mutated without invalidating the signature.

    This is what :func:`sign_commit_ed25519` signs — not the bare
    ``commit_id``.  The verifier must recompute this payload from the stored
    record fields and then call :func:`verify_commit_ed25519`.

    ``branch``, ``repo_id``, and ``metadata`` are intentionally excluded:
    they are mutable by design (a commit is reachable from multiple branches
    after a merge) and their mutation does not represent an integrity
    violation.

    Null bytes are used as field separators to prevent injection attacks
    from field values that contain the separator character.

    Canonical payload format::

        muse-provenance-v2\\n
        <commit_id>\\x00<author>\\x00<agent_id>\\x00<model_id>\\x00<toolchain_id>
        \\x00<prompt_hash>\\x00<committed_at>

    Args:
        commit_id:     SHA-256 hex commit ID (canonical content identity).
        author:        Display author name / email.
        agent_id:      Stable agent identifier.
        model_id:      Model name/version (empty for humans).
        toolchain_id:  Toolchain producing the commit.
        prompt_hash:   SHA-256 hex of the instruction prompt (privacy-preserving).
        committed_at:  ISO-8601 timestamp of the commit (e.g. ``"2026-04-08T12:00:00"``).
                       Empty string is accepted for backward compatibility with unsigned commits.

    Returns:
        64-character lowercase hex SHA-256 digest of the combined payload.
    """
    fields = [commit_id, author, agent_id, model_id, toolchain_id, prompt_hash, committed_at]
    raw = (_PROV_VERSION_PREFIX + _PROV_SEP.join(fields)).encode()
    return hashlib.sha256(raw).hexdigest()


# ---------------------------------------------------------------------------
# Agent identity
# ---------------------------------------------------------------------------


class AgentIdentity(TypedDict, total=False):
    """Structured identity record for a human or AI agent.

    All fields are optional so that partial provenance (e.g. only
    ``agent_id`` is known) can be expressed without filling dummy values.

    ``agent_id``
        Stable human-readable identifier chosen by the agent or its operator.
        Should be unique within a team (e.g. ``"counterpoint-bot-v1"``).
    ``model_id``
        Model identifier for AI agents (e.g. ``"claude-opus-4"``).
        Empty for human authors.
    ``toolchain_id``
        Build system or IDE that produced the commit
        (e.g. ``"cursor-agent-v2"``).
    ``prompt_hash``
        SHA-256 hex of the instruction/prompt that triggered this session.
        Privacy-preserving: the hash is logged without storing the content.
    ``execution_context_hash``
        SHA-256 hex of any additional execution context (system prompt,
        environment config, etc.).
    """

    agent_id: str
    model_id: str
    toolchain_id: str
    prompt_hash: str
    execution_context_hash: str


def make_agent_identity(
    agent_id: str,
    *,
    model_id: str = "",
    toolchain_id: str = "",
    prompt: str = "",
    execution_context: str = "",
) -> AgentIdentity:
    """Build an :class:`AgentIdentity` with optional hashed sensitive fields.

    ``prompt`` and ``execution_context`` are hashed before storage so that
    the raw instruction text never appears in the commit record.

    Args:
        agent_id:           Stable agent identifier string.
        model_id:           Model name/version (empty for humans).
        toolchain_id:       Toolchain producing the commit.
        prompt:             Raw instruction text to hash (not stored).
        execution_context:  Additional context to hash (not stored).

    Returns:
        An :class:`AgentIdentity` with only non-empty fields populated.
    """
    identity = AgentIdentity(agent_id=agent_id)
    if model_id:
        identity["model_id"] = model_id
    if toolchain_id:
        identity["toolchain_id"] = toolchain_id
    if prompt:
        identity["prompt_hash"] = hashlib.sha256(prompt.encode()).hexdigest()
    if execution_context:
        identity["execution_context_hash"] = hashlib.sha256(
            execution_context.encode()
        ).hexdigest()
    return identity


# ---------------------------------------------------------------------------
# Ed25519 signing and verification
# ---------------------------------------------------------------------------


def sign_commit_ed25519(payload: str, private_key: Ed25519PrivateKey) -> str:
    """Sign *payload* with an Ed25519 private key.

    Args:
        payload:     Hex SHA-256 provenance payload from :func:`provenance_payload`.
        private_key: ``Ed25519PrivateKey`` instance from the ``cryptography`` package.

    Returns:
        Base64url-encoded signature (no padding), 88 characters.
    """
    sig_bytes = private_key.sign(payload.encode())
    return base64.urlsafe_b64encode(sig_bytes).rstrip(b"=").decode("ascii")


def verify_commit_ed25519(payload: str, signature_b64: str, public_key_bytes: bytes) -> bool:
    """Verify an Ed25519 *signature_b64* over *payload* using *public_key_bytes*.

    Args:
        payload:          Hex SHA-256 provenance payload from :func:`provenance_payload`.
        signature_b64:    Base64url-encoded signature (with or without padding).
        public_key_bytes: Raw 32-byte Ed25519 public key.

    Returns:
        ``True`` when the signature is valid, ``False`` otherwise.
    """
    from cryptography.exceptions import InvalidSignature
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey

    # Restore padding if stripped.
    padded = signature_b64 + "=" * (-len(signature_b64) % 4)
    try:
        sig_bytes = base64.urlsafe_b64decode(padded)
    except Exception:
        return False

    try:
        pub_key = Ed25519PublicKey.from_public_bytes(public_key_bytes)
        pub_key.verify(sig_bytes, payload.encode())
        return True
    except InvalidSignature:
        return False
    except Exception:
        return False


def public_key_fingerprint(public_key_bytes: bytes) -> str:
    """Return the first 16 hex characters of SHA-256(*public_key_bytes*).

    Short enough to log, long enough for practical uniqueness.

    Args:
        public_key_bytes: Raw 32-byte Ed25519 public key.

    Returns:
        16-character lowercase hex string.
    """
    return hashlib.sha256(public_key_bytes).hexdigest()[:16]


def encode_public_key(private_key: Ed25519PrivateKey) -> tuple[bytes, str]:
    """Extract and encode the public key from an Ed25519 private key.

    Args:
        private_key: ``Ed25519PrivateKey`` instance.

    Returns:
        ``(raw_bytes, b64url_no_padding)`` — the 32-byte raw public key and
        its base64url encoding (no padding, 43 characters).
    """
    from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat

    pub_key = private_key.public_key()
    raw_bytes = pub_key.public_bytes(Encoding.Raw, PublicFormat.Raw)
    b64 = base64.urlsafe_b64encode(raw_bytes).rstrip(b"=").decode("ascii")
    return raw_bytes, b64


# ---------------------------------------------------------------------------
# Convenience: sign a CommitRecord in-place
# ---------------------------------------------------------------------------


def sign_commit_record(
    commit_id: str,
    agent_id: str,
    private_key: Ed25519PrivateKey,
    *,
    author: str = "",
    model_id: str = "",
    toolchain_id: str = "",
    prompt_hash: str = "",
    committed_at: str = "",
) -> tuple[str, str, str] | None:
    """Sign the provenance payload (v2) for *commit_id* with *private_key*.

    Computes :func:`provenance_payload` from the supplied fields so the
    signature covers both the content identity (``commit_id``) and the
    authorship claims, including the ``committed_at`` timestamp.

    Args:
        commit_id:    SHA-256 hex commit ID.
        agent_id:     Stable agent identifier (metadata only; key is from hub identity).
        private_key:  ``Ed25519PrivateKey`` instance.
        author:       Display author name / email.
        model_id:     Model name/version (empty for humans).
        toolchain_id: Toolchain producing the commit.
        prompt_hash:  SHA-256 hex of the instruction prompt.
        committed_at: ISO-8601 timestamp — binds the wall-clock time to the signature.

    Returns:
        ``(signature_b64, public_key_b64, key_fingerprint)`` on success.
    """
    payload = provenance_payload(
        commit_id,
        author=author,
        agent_id=agent_id,
        model_id=model_id,
        toolchain_id=toolchain_id,
        prompt_hash=prompt_hash,
        committed_at=committed_at,
    )
    sig = sign_commit_ed25519(payload, private_key)
    raw_bytes, pub_b64 = encode_public_key(private_key)
    fprint = public_key_fingerprint(raw_bytes)
    logger.debug("✅ Ed25519-signed commit %s with key %s", commit_id[:8], fprint)
    return sig, pub_b64, fprint