"""Tests for agent-first signing — compound identity keys, MUSE_AGENT_KEY injection,
and end-to-end agent commit signing.

Covers:
- identity.py: compound key load/save/clear, provisioned_by field
- keypair.py: agent-specific key paths and generation
- config.py: MUSE_AGENT_KEY env var, agent_id resolution order
- commit.py: signing uses agent key when --agent-id is set
- resolve_signing_identity: agent key → human key fallback
"""
from __future__ import annotations

import pathlib
import json

import pytest

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

from muse.core.identity import (
    IdentityEntry,
    _identity_key,
    clear_identity,
    hostname_from_url,
    list_all_identities,
    load_identity,
    resolve_signing_identity,
    save_identity,
)
from muse.core.keypair import (
    generate_keypair,
    key_path_for,
    load_private_key,
    load_private_key_from_pem,
)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture()
def isolated_identity(tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch) -> pathlib.Path:
    """Redirect the identity store to a temp directory for test isolation."""
    import muse.core.identity as _id_mod
    identity_dir = tmp_path / ".muse"
    identity_dir.mkdir()
    monkeypatch.setattr(_id_mod, "_IDENTITY_DIR", identity_dir)
    monkeypatch.setattr(_id_mod, "_IDENTITY_FILE", identity_dir / "identity.toml")
    return identity_dir


@pytest.fixture()
def isolated_keys(tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch) -> pathlib.Path:
    """Redirect key storage to a temp directory."""
    import muse.core.keypair as _kp_mod
    keys_dir = tmp_path / ".muse" / "keys"
    keys_dir.mkdir(parents=True)
    monkeypatch.setattr(_kp_mod, "_KEYS_DIR", keys_dir)
    return keys_dir


# ---------------------------------------------------------------------------
# _identity_key helper
# ---------------------------------------------------------------------------


class TestIdentityKey:
    def test_human_key_is_bare_hostname(self) -> None:
        assert _identity_key("localhost:10003") == "localhost:10003"

    def test_agent_key_uses_hash_separator(self) -> None:
        assert _identity_key("localhost:10003", "agent-abc") == "localhost:10003#agent-abc"

    def test_none_agent_id_gives_bare_hostname(self) -> None:
        assert _identity_key("musehub.ai", None) == "musehub.ai"

    def test_empty_agent_id_gives_bare_hostname(self) -> None:
        # empty string is falsy
        assert _identity_key("musehub.ai", "") == "musehub.ai"


# ---------------------------------------------------------------------------
# load_identity / save_identity compound keys
# ---------------------------------------------------------------------------


class TestCompoundIdentityKeys:
    def test_human_and_agent_entries_coexist(
        self, isolated_identity: pathlib.Path
    ) -> None:
        hub = "http://localhost:10003"

        human: IdentityEntry = {
            "type": "human",
            "handle": "gabriel",
            "key_path": "/fake/human.pem",
            "algorithm": "ed25519",
            "fingerprint": "a" * 64,
        }
        agent: IdentityEntry = {
            "type": "agent",
            "handle": "agentception-abc",
            "key_path": "/fake/agent.pem",
            "algorithm": "ed25519",
            "fingerprint": "b" * 64,
            "provisioned_by": "gabriel",
        }

        save_identity(hub, human)
        save_identity(hub, agent, agent_id="agentception-abc")

        loaded_human = load_identity(hub)
        loaded_agent = load_identity(hub, agent_id="agentception-abc")

        assert loaded_human is not None
        assert loaded_human["handle"] == "gabriel"
        assert loaded_human["type"] == "human"

        assert loaded_agent is not None
        assert loaded_agent["handle"] == "agentception-abc"
        assert loaded_agent["type"] == "agent"
        assert loaded_agent.get("provisioned_by") == "gabriel"

    def test_agent_entry_does_not_shadow_human(
        self, isolated_identity: pathlib.Path
    ) -> None:
        hub = "http://localhost:10003"
        human: IdentityEntry = {"type": "human", "handle": "gabriel"}
        save_identity(hub, human)

        # Load without agent_id → human entry
        loaded = load_identity(hub)
        assert loaded is not None
        assert loaded["handle"] == "gabriel"

    def test_load_missing_agent_returns_none(
        self, isolated_identity: pathlib.Path
    ) -> None:
        hub = "http://localhost:10003"
        assert load_identity(hub, agent_id="nonexistent") is None

    def test_clear_agent_identity_leaves_human_intact(
        self, isolated_identity: pathlib.Path
    ) -> None:
        hub = "http://localhost:10003"
        human: IdentityEntry = {"type": "human", "handle": "gabriel"}
        agent: IdentityEntry = {"type": "agent", "handle": "agentception-abc", "provisioned_by": "gabriel"}

        save_identity(hub, human)
        save_identity(hub, agent, agent_id="agentception-abc")

        cleared = clear_identity(hub, agent_id="agentception-abc")
        assert cleared is True

        assert load_identity(hub) is not None        # human still present
        assert load_identity(hub, agent_id="agentception-abc") is None  # agent gone

    def test_provisioned_by_roundtrip(
        self, isolated_identity: pathlib.Path
    ) -> None:
        """provisioned_by survives a save/load cycle."""
        hub = "http://localhost:10003"
        agent: IdentityEntry = {
            "type": "agent",
            "handle": "bot-001",
            "provisioned_by": "gabriel",
            "algorithm": "ed25519",
            "fingerprint": "c" * 64,
        }
        save_identity(hub, agent, agent_id="bot-001")
        loaded = load_identity(hub, agent_id="bot-001")
        assert loaded is not None
        assert loaded.get("provisioned_by") == "gabriel"

    def test_list_all_includes_compound_keys(
        self, isolated_identity: pathlib.Path
    ) -> None:
        hub = "http://localhost:10003"
        save_identity(hub, {"type": "human", "handle": "gabriel"})
        save_identity(hub, {"type": "agent", "handle": "bot"}, agent_id="bot")

        all_ids = list_all_identities()
        assert "localhost:10003" in all_ids
        assert "localhost:10003#bot" in all_ids


# ---------------------------------------------------------------------------
# keypair agent-specific paths
# ---------------------------------------------------------------------------


class TestAgentKeyPaths:
    def test_human_key_path(self) -> None:
        p = key_path_for("localhost:10003")
        assert "__" not in p.name
        assert p.name == "localhost_10003.pem"

    def test_agent_key_path_uses_double_underscore(self) -> None:
        p = key_path_for("localhost:10003", "agentception-abc")
        assert p.name == "localhost_10003__agentception-abc.pem"

    def test_generate_agent_keypair_creates_distinct_file(
        self, isolated_keys: pathlib.Path
    ) -> None:
        pub_human, fp_human = generate_keypair("localhost:10003")
        pub_agent, fp_agent = generate_keypair("localhost:10003", "agentception-abc")

        assert fp_human != fp_agent  # different keys
        assert (isolated_keys / "localhost_10003.pem").is_file()
        assert (isolated_keys / "localhost_10003__agentception-abc.pem").is_file()

    def test_load_private_key_with_agent_id(
        self, isolated_keys: pathlib.Path
    ) -> None:
        generate_keypair("localhost:10003", "bot-42")
        key = load_private_key("localhost:10003", "bot-42")
        assert key is not None
        assert isinstance(key, Ed25519PrivateKey)

    def test_load_private_key_agent_id_not_found_returns_none(
        self, isolated_keys: pathlib.Path
    ) -> None:
        # Human key exists but no agent key
        generate_keypair("localhost:10003")
        key = load_private_key("localhost:10003", "nonexistent-agent")
        assert key is None


# ---------------------------------------------------------------------------
# resolve_signing_identity — agent key → fallback chain
# ---------------------------------------------------------------------------


class TestResolveSigningIdentity:
    def test_agent_key_used_when_registered(
        self,
        isolated_identity: pathlib.Path,
        isolated_keys: pathlib.Path,
    ) -> None:
        hub = "http://localhost:10003"
        hostname = hostname_from_url(hub)
        agent_id = "agentception-abc"

        # Generate two distinct keys
        generate_keypair(hostname)
        generate_keypair(hostname, agent_id)

        agent_key_path = str(key_path_for(hostname, agent_id))
        save_identity(
            hub,
            {"type": "agent", "handle": agent_id, "key_path": agent_key_path, "algorithm": "ed25519"},
            agent_id=agent_id,
        )

        result = resolve_signing_identity(hub, agent_id=agent_id)
        assert result is not None
        handle, private_key = result
        assert handle == agent_id
        assert isinstance(private_key, Ed25519PrivateKey)

    def test_falls_back_to_human_when_no_agent_key(
        self,
        isolated_identity: pathlib.Path,
        isolated_keys: pathlib.Path,
    ) -> None:
        hub = "http://localhost:10003"
        hostname = hostname_from_url(hub)

        # Only human key registered
        generate_keypair(hostname)
        human_key_path = str(key_path_for(hostname))
        save_identity(
            hub,
            {"type": "human", "handle": "gabriel", "key_path": human_key_path, "algorithm": "ed25519"},
        )

        # Ask for agent signing but no agent entry → falls back to human
        result = resolve_signing_identity(hub, agent_id="unregistered-agent")
        assert result is not None
        handle, _ = result
        assert handle == "gabriel"

    def test_no_identity_returns_none(
        self, isolated_identity: pathlib.Path
    ) -> None:
        assert resolve_signing_identity("http://localhost:10003") is None

    def test_human_resolve_without_agent_id(
        self,
        isolated_identity: pathlib.Path,
        isolated_keys: pathlib.Path,
    ) -> None:
        hub = "http://localhost:10003"
        hostname = hostname_from_url(hub)
        generate_keypair(hostname)
        save_identity(
            hub,
            {
                "type": "human",
                "handle": "gabriel",
                "key_path": str(key_path_for(hostname)),
                "algorithm": "ed25519",
            },
        )
        result = resolve_signing_identity(hub)
        assert result is not None
        handle, _ = result
        assert handle == "gabriel"


# ---------------------------------------------------------------------------
# load_private_key_from_pem (MUSE_AGENT_KEY env var parsing)
# ---------------------------------------------------------------------------


class TestLoadPrivateKeyFromPem:
    def _make_pem(self) -> bytes:
        key = Ed25519PrivateKey.generate()
        return key.private_bytes(Encoding.PEM, PrivateFormat.PKCS8, NoEncryption())

    def test_valid_pem_returns_key(self) -> None:
        pem = self._make_pem()
        key = load_private_key_from_pem(pem)
        assert isinstance(key, Ed25519PrivateKey)

    def test_invalid_pem_returns_none(self) -> None:
        assert load_private_key_from_pem(b"not a pem") is None

    def test_empty_bytes_returns_none(self) -> None:
        assert load_private_key_from_pem(b"") is None


# ---------------------------------------------------------------------------
# get_signing_identity — MUSE_AGENT_KEY env var
# ---------------------------------------------------------------------------


class TestGetSigningIdentityEnvVar:
    def _make_pem(self) -> str:
        key = Ed25519PrivateKey.generate()
        return key.private_bytes(Encoding.PEM, PrivateFormat.PKCS8, NoEncryption()).decode()

    def test_muse_agent_key_takes_priority(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        from muse.cli.config import get_signing_identity

        pem_str = self._make_pem()
        monkeypatch.setenv("MUSE_AGENT_KEY", pem_str)
        monkeypatch.setenv("MUSE_AGENT_HANDLE", "agentception-xyz")

        # Even with no hub configured, should succeed via env var
        result = get_signing_identity(repo_root=tmp_path, agent_id="agentception-xyz")
        assert result is not None
        assert result.handle == "agentception-xyz"
        assert isinstance(result.private_key, Ed25519PrivateKey)

    def test_muse_agent_key_uses_agent_id_as_fallback_handle(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        from muse.cli.config import get_signing_identity

        pem_str = self._make_pem()
        monkeypatch.setenv("MUSE_AGENT_KEY", pem_str)
        monkeypatch.delenv("MUSE_AGENT_HANDLE", raising=False)

        result = get_signing_identity(repo_root=tmp_path, agent_id="my-agent")
        assert result is not None
        assert result.handle == "my-agent"

    def test_muse_agent_key_default_handle_is_agent(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        from muse.cli.config import get_signing_identity

        pem_str = self._make_pem()
        monkeypatch.setenv("MUSE_AGENT_KEY", pem_str)
        monkeypatch.delenv("MUSE_AGENT_HANDLE", raising=False)

        result = get_signing_identity(repo_root=tmp_path)
        assert result is not None
        assert result.handle == "agent"  # default when no handle provided

    def test_invalid_muse_agent_key_falls_through(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        from muse.cli.config import get_signing_identity

        monkeypatch.setenv("MUSE_AGENT_KEY", "not-valid-pem")

        # Falls through to identity store; no hub configured → None
        result = get_signing_identity(repo_root=tmp_path)
        assert result is None

    def test_no_env_var_and_no_identity_returns_none(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        from muse.cli.config import get_signing_identity

        monkeypatch.delenv("MUSE_AGENT_KEY", raising=False)

        result = get_signing_identity(repo_root=tmp_path)
        assert result is None


# ---------------------------------------------------------------------------
# End-to-end: commit signing uses agent key when --agent-id is set
# ---------------------------------------------------------------------------


class TestCommitSigningWithAgentKey:
    """Verify that muse commit --sign --agent-id uses the agent's dedicated key."""

    def _write_commit_env(
        self,
        tmp_path: pathlib.Path,
        monkeypatch: pytest.MonkeyPatch,
        *,
        agent_pem: str,
        agent_handle: str,
    ) -> "tuple[str, str, str]":
        """Run commit signing logic via MUSE_AGENT_KEY env and return sig/pubkey/key_id."""
        from muse.core.provenance import sign_commit_record

        monkeypatch.setenv("MUSE_AGENT_KEY", agent_pem)
        monkeypatch.setenv("MUSE_AGENT_HANDLE", agent_handle)

        from muse.cli.config import get_signing_identity
        signing = get_signing_identity(repo_root=tmp_path, agent_id=agent_handle)
        assert signing is not None

        result = sign_commit_record(
            "commit-abc123",
            agent_handle,
            signing.private_key,
            author="gabriel",
            model_id="claude-sonnet-4-6",
            toolchain_id="agentception/v1",
            prompt_hash="",
        )
        assert result is not None
        return result

    def test_signing_produces_valid_ed25519_signature(
        self, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
        from cryptography.hazmat.primitives.serialization import (
            Encoding, PrivateFormat, NoEncryption,
        )

        key = Ed25519PrivateKey.generate()
        pem = key.private_bytes(Encoding.PEM, PrivateFormat.PKCS8, NoEncryption()).decode()

        sig, pub_b64, fp = self._write_commit_env(
            tmp_path, monkeypatch, agent_pem=pem, agent_handle="agentception-xyz"
        )
        assert len(sig) == 86  # Ed25519 = 64 bytes → 86 base64url chars
        assert len(pub_b64) == 43  # 32 bytes → 43 base64url chars
        assert len(fp) > 0

    def test_signing_key_fingerprint_matches_public_key(
        self, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        import base64
        import hashlib
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
        from cryptography.hazmat.primitives.serialization import (
            Encoding, PrivateFormat, NoEncryption, PublicFormat,
        )

        key = Ed25519PrivateKey.generate()
        pem = key.private_bytes(Encoding.PEM, PrivateFormat.PKCS8, NoEncryption()).decode()
        pub_raw = key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
        expected_fp = hashlib.sha256(pub_raw).hexdigest()[:16]

        _, _, fp = self._write_commit_env(
            tmp_path, monkeypatch, agent_pem=pem, agent_handle="agent-001"
        )
        # signer_key_id is the first 16 chars of the SHA-256 fingerprint
        assert fp == expected_fp

    def test_agent_signature_verifies_with_embedded_public_key(
        self, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        import base64
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
        from cryptography.hazmat.primitives.serialization import (
            Encoding, PrivateFormat, NoEncryption,
        )
        from muse.core.provenance import provenance_payload, verify_commit_ed25519

        key = Ed25519PrivateKey.generate()
        pem = key.private_bytes(Encoding.PEM, PrivateFormat.PKCS8, NoEncryption()).decode()
        handle = "agentception-verify-test"

        sig, pub_b64, _ = self._write_commit_env(
            tmp_path, monkeypatch, agent_pem=pem, agent_handle=handle
        )

        # Reconstruct the signed payload the same way sign_commit_record does
        payload = provenance_payload(
            "commit-abc123",
            author="gabriel",
            agent_id=handle,
            model_id="claude-sonnet-4-6",
            toolchain_id="agentception/v1",
            prompt_hash="",
        )

        # Decode the embedded public key and verify
        padded = pub_b64 + "=" * (4 - len(pub_b64) % 4)
        pub_bytes = base64.urlsafe_b64decode(padded)
        assert verify_commit_ed25519(payload, sig, pub_bytes)