"""Global identity store — ``~/.muse/identity.toml``.

Ed25519 key-pair credentials are kept here, separate from per-repository
configuration.  This means keys are never accidentally committed to
version control, and a single identity can authenticate across all
repositories on the same hub.

Why global, not per-repo
-------------------------
Muse makes identity a first-class, machine-scoped concept.
The repository knows *where* the hub is (``[hub] url`` in config.toml).
The machine knows *who you are* (this file).  The two concerns are
deliberately separated.

Identity types
--------------
``type = "human"``
    A person.  Authenticated via Ed25519 key pair.

``type = "agent"``
    An autonomous process.  The ``capabilities`` field reflects what the
    agent is allowed to do, enabling self-inspection before attempting an
    operation.  The ``provisioned_by`` field records the human handle that
    authorised this agent, establishing the trust chain at provisioning time.

File format
-----------
TOML with one section per hub hostname or compound ``hostname#agent_id`` key::

    ["localhost:10003"]
    type        = "human"
    handle      = "gabriel"
    key_path    = "/Users/gabriel/.muse/keys/localhost:10003.pem"
    algorithm   = "ed25519"
    fingerprint = "<sha256hex>"

    ["localhost:10003#agentception-abc123"]
    type          = "agent"
    handle        = "agentception-abc123"
    key_path      = "/Users/gabriel/.muse/keys/localhost:10003__agentception-abc123.pem"
    algorithm     = "ed25519"
    fingerprint   = "<sha256hex>"
    provisioned_by = "gabriel"
    capabilities  = ["push", "pull"]

Compound key format
-------------------
``"hostname#agent_id"`` — the ``#`` separator is not a valid hostname
character so it unambiguously separates the hub from the agent handle.
Human entries use the bare hostname; agent entries append ``#<agent_id>``.

Security model
--------------
- ``~/.muse/`` is created with mode 0o700 (user-only directory).
- ``~/.muse/identity.toml`` is written with mode 0o600 **from the first
  byte** — using ``os.open()`` + ``os.fchmod()`` 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.  A kill signal during write
  leaves the old file intact, never a partial file.
- Symlink guard: if the target path is already a symlink, write is refused.
  This blocks symlink-based credential-overwrite attacks.
- The file is never read or written as part of a repository snapshot.
"""

from __future__ import annotations

import contextlib
import fcntl
import logging
import os
import pathlib
import stat
import tempfile
import tomllib
from collections.abc import Generator
from typing import TypedDict

logger = logging.getLogger(__name__)

_IDENTITY_DIR = pathlib.Path.home() / ".muse"
_IDENTITY_FILE = _IDENTITY_DIR / "identity.toml"

type _IdentityMap = dict[str, "IdentityEntry"]


# ---------------------------------------------------------------------------
# Types
# ---------------------------------------------------------------------------


class IdentityEntry(TypedDict, total=False):
    """One authenticated identity, keyed by hub hostname in identity.toml.

    Human entry (``~/.muse/identity.toml``)::

        ["localhost:10003"]
        type        = "human"
        handle      = "gabriel"
        key_path    = "/Users/gabriel/.muse/keys/localhost:10003.pem"
        algorithm   = "ed25519"
        fingerprint = "<sha256hex>"

    Agent entry (compound key ``hostname#agent_id``)::

        ["localhost:10003#agentception-abc123"]
        type           = "agent"
        handle         = "agentception-abc123"
        key_path       = "/Users/gabriel/.muse/keys/localhost:10003__agentception-abc123.pem"
        algorithm      = "ed25519"
        fingerprint    = "<sha256hex>"
        provisioned_by = "gabriel"
        capabilities   = ["push", "pull"]
    """

    type: str             # "human" | "agent"
    handle: str           # hub-assigned handle, e.g. "gabriel" or "agentception-abc123"
    key_path: str         # absolute path to Ed25519 private key PEM
    algorithm: str        # "ed25519" | "ml-dsa-65"
    fingerprint: str      # SHA-256 hex of public key (for display / verification)
    capabilities: list[str]  # agent capability strings (empty for humans)
    provisioned_by: str   # for agents: handle of the human who provisioned this key


# ---------------------------------------------------------------------------
# Path helper
# ---------------------------------------------------------------------------


def get_identity_path() -> pathlib.Path:
    """Return the path to the global identity file (``~/.muse/identity.toml``)."""
    return _IDENTITY_FILE


def _identity_key(hostname: str, agent_id: str | None = None) -> str:
    """Build the TOML section key for a given hostname and optional agent_id.

    Human identities:  ``"localhost:10003"``
    Agent identities:  ``"localhost:10003#agentception-abc123"``

    The ``#`` character is not valid in a hostname so it unambiguously marks
    the boundary between hub and agent handle.
    """
    if agent_id:
        return f"{hostname}#{agent_id}"
    return hostname


# ---------------------------------------------------------------------------
# URL → hostname normalisation
# ---------------------------------------------------------------------------


def hostname_from_url(url: str) -> str:
    """Normalise *url* to a lowercase hostname suitable for use as a dict key.

    Security properties
    -------------------
    - Strips the scheme (``https://``), so different scheme representations of
      the same host resolve to the same key.
    - Strips userinfo (``user:password@``) — embedded credentials in a URL are
      never stored as part of the hostname key.
    - Normalises to lowercase — DNS is case-insensitive, so ``MUSEHUB.AI``
      and ``musehub.ai`` are the same host and must resolve to the same entry.

    Examples::

        "https://musehub.ai/repos/x"      → "musehub.ai"
        "https://admin:s3cr3t@musehub.ai" → "musehub.ai"
        "MUSEHUB.AI"                       → "musehub.ai"
        "https://musehub.ai"               → "musehub.ai"
        "musehub.ai:8443"                  → "musehub.ai:8443"
    """
    stripped = url.strip().rstrip("/")
    # Remove scheme.
    if "://" in stripped:
        stripped = stripped.split("://", 1)[1]
    # Remove userinfo (user:password@) — never embed credentials in the key.
    if "@" in stripped:
        stripped = stripped.rsplit("@", 1)[1]
    # Keep only host[:port], strip any path.
    hostname = stripped.split("/")[0]
    # Normalise to lowercase — DNS is case-insensitive.
    return hostname.lower()


# ---------------------------------------------------------------------------
# TOML serialiser (write-side — stdlib tomllib is read-only)
# ---------------------------------------------------------------------------


def _toml_escape(value: str) -> str:
    """Escape a string value for embedding in a TOML double-quoted string."""
    return value.replace("\\", "\\\\").replace('"', '\\"')


def _dump_identity(identities: _IdentityMap) -> str:
    """Serialise a hostname → entry mapping to TOML text.

    All hostnames are quoted in the section header so that dotted names
    (e.g. ``musehub.ai``) are treated as literal keys, not nested tables.
    All string values are TOML-escaped to prevent injection.
    """
    lines: list[str] = []
    for hostname in sorted(identities):
        entry = identities[hostname]
        # Always quote the section key — dotted names are literal, not nested.
        lines.append(f'["{_toml_escape(hostname)}"]')
        t = entry.get("type", "")
        if t:
            lines.append(f'type = "{_toml_escape(t)}"')
        handle = entry.get("handle", "")
        if handle:
            lines.append(f'handle = "{_toml_escape(handle)}"')
        key_path = entry.get("key_path", "")
        if key_path:
            lines.append(f'key_path = "{_toml_escape(key_path)}"')
        algorithm = entry.get("algorithm", "")
        if algorithm:
            lines.append(f'algorithm = "{_toml_escape(algorithm)}"')
        fingerprint = entry.get("fingerprint", "")
        if fingerprint:
            lines.append(f'fingerprint = "{_toml_escape(fingerprint)}"')
        provisioned_by = entry.get("provisioned_by", "")
        if provisioned_by:
            lines.append(f'provisioned_by = "{_toml_escape(provisioned_by)}"')
        caps = entry.get("capabilities") or []
        if caps:
            caps_str = ", ".join(f'"{_toml_escape(c)}"' for c in caps)
            lines.append(f"capabilities = [{caps_str}]")
        lines.append("")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Load / save
# ---------------------------------------------------------------------------


def _load_all(path: pathlib.Path) -> _IdentityMap:
    """Load all identity entries from *path*.  Returns empty dict if absent."""
    if not path.is_file():
        return {}
    try:
        with path.open("rb") as fh:
            raw = tomllib.load(fh)
    except Exception as exc:  # noqa: BLE001
        # Log only the exception *type*, never its message — a TOML parse
        # error surfaced by tomllib includes the offending line, which can
        # contain a fragment of the token being written when the file is corrupt.
        logger.warning(
            "⚠️ Failed to parse identity file %s (%s — run `muse auth register` to re-authenticate)",
            path,
            type(exc).__name__,
        )
        return {}

    result: _IdentityMap = {}
    for hostname, raw_entry in raw.items():
        if not isinstance(raw_entry, dict):
            continue
        entry: IdentityEntry = {}
        t = raw_entry.get("type")
        if isinstance(t, str):
            entry["type"] = t
        handle = raw_entry.get("handle")
        if isinstance(handle, str):
            entry["handle"] = handle
        key_path = raw_entry.get("key_path")
        if isinstance(key_path, str):
            entry["key_path"] = key_path
        algorithm = raw_entry.get("algorithm")
        if isinstance(algorithm, str):
            entry["algorithm"] = algorithm
        fingerprint = raw_entry.get("fingerprint")
        if isinstance(fingerprint, str):
            entry["fingerprint"] = fingerprint
        caps = raw_entry.get("capabilities")
        if isinstance(caps, list):
            entry["capabilities"] = [str(c) for c in caps if isinstance(c, str)]
        provisioned_by = raw_entry.get("provisioned_by")
        if isinstance(provisioned_by, str):
            entry["provisioned_by"] = provisioned_by
        result[hostname] = entry

    return result


@contextlib.contextmanager
def _identity_write_lock() -> Generator[None, None, None]:
    """Acquire an exclusive advisory write-lock on the identity store.

    Uses a dedicated lock file (``~/.muse/.identity.lock``) so that the lock
    survives the atomic rename of ``identity.toml`` itself.

    Advisory (cooperative) locking protects all Muse processes that use this
    lock against concurrent read-modify-write races.  Direct file edits by
    external tools bypass the lock — that is acceptable; the user is then
    responsible for data consistency.

    POSIX-only (``fcntl.flock``).  The lock is blocking with no timeout;
    CLI commands are short-lived and lock contention is expected to be brief.
    """
    lock_path = _IDENTITY_DIR / ".identity.lock"
    _IDENTITY_DIR.mkdir(parents=True, exist_ok=True)
    # Create the lock file with owner-only permissions; O_CLOEXEC prevents
    # child processes from inheriting the file descriptor.
    lock_fd = os.open(
        str(lock_path),
        os.O_CREAT | os.O_WRONLY | os.O_CLOEXEC,
        stat.S_IRUSR | stat.S_IWUSR,
    )
    try:
        fcntl.flock(lock_fd, fcntl.LOCK_EX)
        try:
            yield
        finally:
            fcntl.flock(lock_fd, fcntl.LOCK_UN)
    finally:
        os.close(lock_fd)


def _save_all(identities: _IdentityMap, path: pathlib.Path) -> None:
    """Write *identities* to *path* securely.

    Security guarantees
    -------------------
    1. **Symlink guard** — refuses to write if *path* is already a symlink,
       preventing an attacker from pre-placing a symlink to a file they want
       overwritten.
    2. **0o700 directory** — ``~/.muse/`` is restricted to the owner so other
       local users cannot list or traverse it.
    3. **0o600 from byte zero** — the temp file is ``fchmod``-ed to 0o600
       *before* any data is written, eliminating the TOCTOU window that
       ``write_text()`` + ``chmod()`` creates.
    4. **Atomic rename** — ``os.replace()`` swaps the temp file over the
       target atomically; a kill signal during write leaves the old file intact.
    """
    dir_path = path.parent

    # 1. Create ~/.muse/ with owner-only permissions (0o700).
    dir_path.mkdir(parents=True, exist_ok=True)
    try:
        os.chmod(dir_path, stat.S_IRWXU)  # 0o700
    except OSError as exc:
        logger.warning("⚠️ Could not set permissions on %s: %s", dir_path, exc)

    # 2. Symlink guard — never follow a symlink placed at the target path.
    if path.is_symlink():
        raise OSError(
            f"Security: {path} is a symlink. "
            "Refusing to write credentials to a symlink target."
        )

    text = _dump_identity(identities)

    # 3. Write to a temp file in the same directory (same fs → atomic rename).
    #    Set 0o600 via fchmod *before* writing any data.
    fd, tmp_path_str = tempfile.mkstemp(dir=dir_path, prefix=".identity-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, "w", encoding="utf-8") as fh:
            fh.write(text)
        # 4. Atomic rename — old file stays intact if we crash before this.
        os.replace(tmp_path, path)
    except Exception:
        try:
            tmp_path.unlink(missing_ok=True)
        except OSError:
            pass
        raise


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


def load_identity(hub_url: str, agent_id: str | None = None) -> IdentityEntry | None:
    """Return the stored identity for *hub_url* (and optional *agent_id*), or ``None``.

    The URL is normalised to a hostname before lookup, so
    ``https://musehub.ai/repos/x`` and ``musehub.ai`` resolve to the same
    entry.

    Args:
        hub_url:  Hub URL or bare hostname.
        agent_id: Agent handle (e.g. ``"agentception-abc123"``).  When
                  supplied the lookup key is ``"hostname#agent_id"``,
                  i.e. the agent's dedicated entry rather than the human's.

    Returns:
        :class:`IdentityEntry` if an identity is stored, else ``None``.
    """
    hostname = hostname_from_url(hub_url)
    key = _identity_key(hostname, agent_id)
    return _load_all(_IDENTITY_FILE).get(key)


def save_identity(hub_url: str, entry: IdentityEntry, agent_id: str | None = None) -> None:
    """Store *entry* as the identity for *hub_url* (and optional *agent_id*).

    The entire read-modify-write cycle is wrapped in an exclusive advisory
    lock so that concurrent ``muse auth register`` calls (e.g. from parallel
    agents) cannot race and overwrite each other's entries.

    Creates ``~/.muse/identity.toml`` with mode 0o600 if it does not exist.

    Args:
        hub_url:  Hub URL or bare hostname.
        entry:    Identity data to store.
        agent_id: When supplied, the entry is stored under the compound key
                  ``"hostname#agent_id"`` rather than the bare hostname.
    """
    hostname = hostname_from_url(hub_url)
    key = _identity_key(hostname, agent_id)
    with _identity_write_lock():
        identities = _load_all(_IDENTITY_FILE)
        identities[key] = entry
        _save_all(identities, _IDENTITY_FILE)
    logger.info("✅ Identity for %s saved", key)


def clear_identity(hub_url: str, agent_id: str | None = None) -> bool:
    """Remove the stored identity for *hub_url* (and optional *agent_id*).

    The entire read-modify-write cycle is wrapped in an exclusive advisory
    lock (see :func:`save_identity`).

    Args:
        hub_url:  Hub URL or bare hostname.
        agent_id: When supplied, removes the agent's compound-key entry.

    Returns:
        ``True`` if an entry was removed, ``False`` if no entry existed.
    """
    hostname = hostname_from_url(hub_url)
    key = _identity_key(hostname, agent_id)
    with _identity_write_lock():
        identities = _load_all(_IDENTITY_FILE)
        if key not in identities:
            return False
        del identities[key]
        _save_all(identities, _IDENTITY_FILE)
    logger.info("✅ Identity for %s cleared", key)
    return True


def _check_key_file_permissions(key_path: pathlib.Path) -> bool:
    """Return ``True`` when *key_path* has safe permissions (owner-only, 0o600 or stricter).

    Checks:
    1. File mode must not expose group-read (0o040) or world-read (0o004) bits.
    2. File must be owned by the current user (skipped when uid == 0).

    Logs a WARNING and returns ``False`` when either check fails, including
    the fix command so the user can resolve the problem.

    Args:
        key_path: Path to the private key PEM file.

    Returns:
        ``True`` when permissions are safe, ``False`` otherwise.
    """
    try:
        st = key_path.stat()
    except OSError as exc:
        logger.warning("⚠️ Cannot stat key file %s: %s", key_path, exc)
        return False

    # Permissions: any bit beyond owner read/write (0o600) is unsafe.
    # 0o177 masks all bits that are NOT owner read/write:
    #   0o177 = 0o177 (all group, world, and special bits)
    mode_octal = st.st_mode & 0o777
    if mode_octal & 0o177:
        logger.warning(
            "⚠️ Key file %s has permissions %04o — expected 0600 or stricter. "
            "Group/world-readable private key files are refused for security. "
            "Fix with: chmod 600 %s",
            key_path,
            mode_octal,
            key_path,
        )
        return False

    # Ownership: refuse keys owned by another user (root is exempt).
    current_uid = os.getuid()
    if current_uid != 0 and st.st_uid != current_uid:
        logger.warning(
            "⚠️ Key file %s is owned by UID %d but running as UID %d — "
            "refusing to use a key owned by another user. "
            "Ensure the key is owned by you, or use a different key path.",
            key_path,
            st.st_uid,
            current_uid,
        )
        return False

    return True


def _load_private_key_from_path(key_path_str: str) -> "object | None":
    """Load an Ed25519 private key from *key_path_str*.  Returns ``None`` on any failure.

    Permission checks performed before loading:
    - File mode must be 0o600 or stricter (no group/world read bits).
    - File must be owned by the current user (root is exempt from ownership check).

    Returns ``None`` (with a WARNING) when either check fails, preventing
    accidental use of an insecure or foreign key.
    """
    key_path = pathlib.Path(key_path_str)
    if not key_path.is_file():
        logger.warning("⚠️ Key file not found: %s", key_path)
        return None

    # Security gate: refuse keys with unsafe permissions or foreign ownership.
    if not _check_key_file_permissions(key_path):
        return None

    try:
        from cryptography.hazmat.primitives.serialization import load_pem_private_key
        pem_bytes = key_path.read_bytes()
        return load_pem_private_key(pem_bytes, password=None)
    except Exception as exc:
        logger.warning("⚠️ Could not load private key from %s: %s", key_path, exc)
        return None


def resolve_signing_identity(
    hub_url: str,
    agent_id: str | None = None,
) -> "tuple[str, object] | None":
    """Return ``(handle, private_key)`` for *hub_url*, or ``None``.

    Resolution order when *agent_id* is provided:
    1. Agent-specific entry (``"hostname#agent_id"``) — used if present.
    2. Human entry (bare hostname) — fallback for operators who have not yet
       registered a dedicated agent key.

    When *agent_id* is ``None``, only the human (bare-hostname) entry is tried.

    The private key is loaded from the ``key_path`` stored in ``identity.toml``.
    Returns ``None`` when no identity is configured or the key file is missing.

    Args:
        hub_url:  Hub URL or bare hostname.
        agent_id: Optional agent handle to look up a dedicated agent key first.

    Returns:
        ``(handle, Ed25519PrivateKey)`` tuple, or ``None``.
    """
    # 1. Try agent-specific entry when an agent_id is given.
    if agent_id:
        agent_entry = load_identity(hub_url, agent_id=agent_id)
        if agent_entry is not None:
            handle = agent_entry.get("handle", "")
            key_path_str = agent_entry.get("key_path", "")
            if handle and key_path_str:
                private_key = _load_private_key_from_path(key_path_str)
                if private_key is not None:
                    logger.debug("✅ Agent signing key resolved: handle=%s", handle)
                    return handle, private_key

    # 2. Fall back to the human entry.
    human_entry = load_identity(hub_url)
    if human_entry is None:
        return None

    handle = human_entry.get("handle", "")
    key_path_str = human_entry.get("key_path", "")

    if not handle or not key_path_str:
        return None

    private_key = _load_private_key_from_path(key_path_str)
    if private_key is None:
        return None

    if agent_id:
        logger.debug(
            "⚠️ No dedicated key for agent '%s' — signing with human key (handle=%s). "
            "Run `muse auth keygen --agent-id %s` to provision a dedicated agent key.",
            agent_id, handle, agent_id,
        )
    return handle, private_key


def list_all_identities() -> _IdentityMap:
    """Return all stored identities keyed by hub hostname.

    Returns an empty dict if the identity file does not exist.
    """
    return _load_all(_IDENTITY_FILE)


def clear_all_identities() -> list[str]:
    """Remove every stored identity in a single atomic write.

    The entire read-modify-write cycle is wrapped in an exclusive advisory
    lock (see :func:`save_identity`) so no concurrent process can race.

    This is O(1) writes regardless of how many identities are stored —
    far cheaper than calling :func:`clear_identity` in a loop, which does
    N separate lock → read → write cycles.

    Returns:
        Sorted list of hostnames that were removed.  Empty list if the
        identity file did not exist or contained no entries.
    """
    with _identity_write_lock():
        identities = _load_all(_IDENTITY_FILE)
        if not identities:
            return []
        removed = sorted(identities.keys())
        _save_all({}, _IDENTITY_FILE)
    logger.info("✅ All identities cleared (%d hub(s))", len(removed))
    return removed