"""Unified TDD tests for HD-only keygen architecture.

This file validates:
  - agent_id_to_slot: stable, deterministic, BIP32-safe slot mapping
  - Human keygen: fresh mnemonic, no --hd flag needed (HD is the only mode)
  - Agent keygen: derived from operator's mnemonic via derive_agent_sub_seed
  - run_recover: re-derives same fingerprint from stored mnemonic
  - No JBOK: generate_keypair must not exist
  - Integration flow: keygen → agent keygen → recover round-trip
"""

from __future__ import annotations

import base64
import hashlib
import json
import pathlib

import pytest
from cryptography.hazmat.primitives.serialization import load_pem_private_key

from tests.cli_test_helper import CliRunner
from muse.core import keypair as kp_module
from muse.core import identity as id_module
from muse.core.bip39 import mnemonic_to_seed, validate_mnemonic
from muse.core.hdkeys import (
    DOMAIN_IDENTITY,
    ENTITY_AGENT,
    ENTITY_HUMAN,
    MUSE_PURPOSE,
    ROLE_SIGN,
    agent_id_to_slot,
    derive_agent_sub_seed,
    derive_identity_key,
    muse_path,
)

runner = CliRunner()

_HUB = "http://localhost:10003"
_HOSTNAME = "localhost:10003"
# A well-known BIP39 test mnemonic (abandon × 11 + about)
_TEST_MNEMONIC_12 = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _patch_home(monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> pathlib.Path:
    fake_home = tmp_path / "home"
    fake_home.mkdir(parents=True, exist_ok=True)
    monkeypatch.setattr(pathlib.Path, "home", staticmethod(lambda: fake_home))
    monkeypatch.setattr(kp_module, "_KEYS_DIR", fake_home / ".muse" / "keys")
    monkeypatch.setattr(id_module, "_IDENTITY_DIR", fake_home / ".muse")
    monkeypatch.setattr(id_module, "_IDENTITY_FILE", fake_home / ".muse" / "identity.toml")
    return fake_home


def _keygen(monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path,
            extra_args: list[str] | None = None) -> tuple[pathlib.Path, object]:
    """Run ``muse auth keygen --hub <HUB>`` and return (fake_home, result)."""
    fake_home = _patch_home(monkeypatch, tmp_path)
    args = ["auth", "keygen", "--hub", _HUB] + (extra_args or [])
    result = runner.invoke(None, args)
    return fake_home, result


# ---------------------------------------------------------------------------
# agent_id_to_slot — unit tests
# ---------------------------------------------------------------------------


class TestAgentIdToSlot:
    """agent_id_to_slot must map handle strings to stable, valid BIP32 indices."""

    def test_returns_int(self) -> None:
        slot = agent_id_to_slot("my-agent")
        assert isinstance(slot, int)

    def test_in_valid_bip32_range(self) -> None:
        """All slots must be in [0, 2^31 - 1] (hardened offset applied by caller)."""
        for handle in ["alpha", "beta", "gamma-007", "a" * 100]:
            slot = agent_id_to_slot(handle)
            assert 0 <= slot <= 0x7FFF_FFFF, f"slot={slot} out of range for {handle!r}"

    def test_deterministic(self) -> None:
        """Same handle must always produce the same slot."""
        handle = "agentception-abc123"
        assert agent_id_to_slot(handle) == agent_id_to_slot(handle)

    def test_distinct_handles_likely_distinct_slots(self) -> None:
        """Different handles should not collide (SHA-256 collision resistance)."""
        handles = ["alice", "bob", "carol", "dave", "eve", "frank"]
        slots = [agent_id_to_slot(h) for h in handles]
        assert len(set(slots)) == len(slots), f"Unexpected slot collision: {slots}"

    def test_known_vector(self) -> None:
        """Verify the slot for 'agentception' against a manually computed value."""
        import hashlib as _hashlib
        handle = "agentception"
        digest = _hashlib.sha256(handle.encode()).digest()
        expected = int.from_bytes(digest[:4], "big") & 0x7FFF_FFFF
        assert agent_id_to_slot(handle) == expected

    def test_empty_string_handled(self) -> None:
        """Edge case: empty string handle should not crash."""
        slot = agent_id_to_slot("")
        assert 0 <= slot <= 0x7FFF_FFFF

    def test_unicode_handle(self) -> None:
        """Unicode agent handles should produce valid slots."""
        slot = agent_id_to_slot("音楽エージェント")
        assert 0 <= slot <= 0x7FFF_FFFF


# ---------------------------------------------------------------------------
# No JBOK — generate_keypair must not exist
# ---------------------------------------------------------------------------


class TestNoJbok:
    """JBOK mode is deleted. generate_keypair must not exist anywhere."""

    def test_generate_keypair_not_in_module(self) -> None:
        import importlib
        kp = importlib.import_module("muse.core.keypair")
        assert not hasattr(kp, "generate_keypair"), \
            "generate_keypair still exists — JBOK was not fully removed"

    def test_generate_keypair_not_importable(self) -> None:
        with pytest.raises(ImportError):
            from muse.core.keypair import generate_keypair  # noqa: F401

    def test_generate_hd_keypair_exists(self) -> None:
        from muse.core.keypair import generate_hd_keypair
        assert callable(generate_hd_keypair)


# ---------------------------------------------------------------------------
# Human keygen — no --hd flag, 24-word default
# ---------------------------------------------------------------------------


class TestHumanKeygen:
    """Human keygen: HD is the only mode. No --hd flag required."""

    def test_exits_zero(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        _, result = _keygen(monkeypatch, tmp_path)
        assert result.exit_code == 0, result.output

    def test_pem_written(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        fake_home, result = _keygen(monkeypatch, tmp_path)
        pem = fake_home / ".muse" / "keys" / "localhost_10003.pem"
        assert pem.is_file(), f"PEM not created. Output:\n{result.output}"

    def test_pem_mode_600(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        fake_home, result = _keygen(monkeypatch, tmp_path)
        pem = fake_home / ".muse" / "keys" / "localhost_10003.pem"
        assert result.exit_code == 0
        mode = pem.stat().st_mode & 0o777
        assert mode == 0o600, f"PEM mode is {oct(mode)}, expected 0o600"

    def test_default_24_word_mnemonic(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        """Default strength=256 produces a 24-word mnemonic (visible on a TTY)."""
        import muse.cli.commands.auth as auth_mod
        monkeypatch.setattr(auth_mod, "_stderr_isatty", lambda: True)
        _, result = _keygen(monkeypatch, tmp_path)
        assert result.exit_code == 0
        all_text = result.output
        mnemonic_line = None
        for line in all_text.splitlines():
            words = line.strip().split()
            if len(words) == 24 and all(w.isalpha() for w in words):
                mnemonic_line = line.strip()
                break
        assert mnemonic_line is not None, f"No 24-word line found:\n{all_text}"
        assert validate_mnemonic(mnemonic_line)

    def test_json_no_mnemonic_in_stdout(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        """Mnemonic is sensitive — must never appear in JSON stdout."""
        _, result = _keygen(monkeypatch, tmp_path, ["--json"])
        payload = json.loads(result.output.splitlines()[0])
        assert "mnemonic" not in payload

    def test_json_mnemonic_word_count_24(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        _, result = _keygen(monkeypatch, tmp_path, ["--json"])
        payload = json.loads(result.output.splitlines()[0])
        assert payload.get("mnemonic_word_count") == 24

    def test_identity_toml_has_no_key_source(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        import tomllib
        fake_home, result = _keygen(monkeypatch, tmp_path)
        assert result.exit_code == 0
        data = tomllib.loads((fake_home / ".muse" / "identity.toml").read_text())
        assert "key_source" not in data[_HOSTNAME]

    def test_force_overwrites_existing(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        _keygen(monkeypatch, tmp_path)
        _, result = _keygen(monkeypatch, tmp_path, ["--force"])
        assert result.exit_code == 0

    def test_no_force_rejects_existing(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        _keygen(monkeypatch, tmp_path)
        _, result = _keygen(monkeypatch, tmp_path)  # second time, no --force
        assert result.exit_code != 0

    def test_strength_128_gives_12_words(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        _, result = _keygen(monkeypatch, tmp_path, ["--strength", "128", "--json"])
        assert result.exit_code == 0
        payload = json.loads(result.output.splitlines()[0])
        assert payload["mnemonic_word_count"] == 12


# ---------------------------------------------------------------------------
# Agent keygen — derived from operator's mnemonic
# ---------------------------------------------------------------------------


class TestAgentKeygen:
    """Agent keys must be derived from the operator's HD mnemonic."""

    def _setup_operator(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> pathlib.Path:
        """Generate a human (operator) key first."""
        fake_home = _patch_home(monkeypatch, tmp_path)
        result = runner.invoke(None, ["auth", "keygen", "--hub", _HUB])
        assert result.exit_code == 0, f"Operator keygen failed:\n{result.output}"
        return fake_home

    def test_agent_keygen_exits_zero(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        self._setup_operator(monkeypatch, tmp_path)
        result = runner.invoke(None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "bot-alpha"])
        assert result.exit_code == 0, result.output

    def test_agent_pem_at_expected_path(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        fake_home = self._setup_operator(monkeypatch, tmp_path)
        runner.invoke(None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "bot-alpha"])
        pem = fake_home / ".muse" / "keys" / "localhost_10003__bot-alpha.pem"
        assert pem.is_file(), f"Agent PEM not found at {pem}"

    def test_agent_pem_mode_600(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        fake_home = self._setup_operator(monkeypatch, tmp_path)
        runner.invoke(None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "bot-alpha"])
        pem = fake_home / ".muse" / "keys" / "localhost_10003__bot-alpha.pem"
        mode = pem.stat().st_mode & 0o777
        assert mode == 0o600

    def test_agent_json_has_hd_path(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        self._setup_operator(monkeypatch, tmp_path)
        result = runner.invoke(
            None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "bot-alpha", "--json"]
        )
        assert result.exit_code == 0, result.output
        payload = json.loads(result.output.splitlines()[0])
        assert "hd_path" in payload
        assert str(MUSE_PURPOSE) in payload["hd_path"]

    def test_agent_json_has_provisioned_by(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        self._setup_operator(monkeypatch, tmp_path)
        result = runner.invoke(
            None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "bot-alpha", "--json"]
        )
        payload = json.loads(result.output.splitlines()[0])
        assert "provisioned_by_fingerprint" in payload
        assert len(payload["provisioned_by_fingerprint"]) == 64  # SHA-256 hex

    def test_agent_key_different_from_human_key(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        fake_home = self._setup_operator(monkeypatch, tmp_path)
        runner.invoke(None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "bot-alpha"])

        human_pem = fake_home / ".muse" / "keys" / "localhost_10003.pem"
        agent_pem = fake_home / ".muse" / "keys" / "localhost_10003__bot-alpha.pem"

        human_key = load_pem_private_key(human_pem.read_bytes(), password=None)
        agent_key = load_pem_private_key(agent_pem.read_bytes(), password=None)

        from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
        human_pub = human_key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
        agent_pub = agent_key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
        assert human_pub != agent_pub, "Agent and human keys must be distinct"

    def test_two_agents_have_distinct_keys(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        fake_home = self._setup_operator(monkeypatch, tmp_path)
        runner.invoke(None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "bot-alpha"])
        runner.invoke(None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "bot-beta"])

        alpha_pem = fake_home / ".muse" / "keys" / "localhost_10003__bot-alpha.pem"
        beta_pem = fake_home / ".muse" / "keys" / "localhost_10003__bot-beta.pem"

        from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
        alpha_key = load_pem_private_key(alpha_pem.read_bytes(), password=None)
        beta_key = load_pem_private_key(beta_pem.read_bytes(), password=None)
        alpha_pub = alpha_key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
        beta_pub = beta_key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
        assert alpha_pub != beta_pub, "Different agent handles must produce different keys"

    def test_agent_keygen_without_operator_exits_nonzero(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        """Attempt to derive agent key before operator key is set up."""
        _patch_home(monkeypatch, tmp_path)
        result = runner.invoke(None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "bot-alpha"])
        assert result.exit_code != 0

    def test_agent_key_deterministic(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        """Same operator mnemonic + same agent handle = same agent key."""
        fake_home = self._setup_operator(monkeypatch, tmp_path)
        result1 = runner.invoke(
            None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "bot-alpha", "--json"]
        )
        fp1 = json.loads(result1.output.splitlines()[0])["fingerprint"]

        # Re-derive: force-overwrite the agent key (same operator mnemonic on disk)
        result2 = runner.invoke(
            None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "bot-alpha", "--force", "--json"]
        )
        fp2 = json.loads(result2.output.splitlines()[0])["fingerprint"]
        assert fp1 == fp2, "Agent key not deterministic given same operator mnemonic + handle"


# ---------------------------------------------------------------------------
# derive_agent_sub_seed — unit tests (no CLI)
# ---------------------------------------------------------------------------


class TestDeriveAgentSubSeed:
    """derive_agent_sub_seed must produce stable, domain-isolated sub-seeds."""

    def test_returns_64_bytes(self) -> None:
        seed = mnemonic_to_seed(_TEST_MNEMONIC_12)
        slot = agent_id_to_slot("bot-alpha")
        sub_seed = derive_agent_sub_seed(seed, DOMAIN_IDENTITY, slot)
        assert len(sub_seed) == 64

    def test_deterministic(self) -> None:
        seed = mnemonic_to_seed(_TEST_MNEMONIC_12)
        slot = agent_id_to_slot("bot-alpha")
        s1 = derive_agent_sub_seed(seed, DOMAIN_IDENTITY, slot)
        s2 = derive_agent_sub_seed(seed, DOMAIN_IDENTITY, slot)
        assert s1 == s2

    def test_different_slots_different_sub_seeds(self) -> None:
        seed = mnemonic_to_seed(_TEST_MNEMONIC_12)
        slot_a = agent_id_to_slot("bot-alpha")
        slot_b = agent_id_to_slot("bot-beta")
        assert slot_a != slot_b
        sub_a = derive_agent_sub_seed(seed, DOMAIN_IDENTITY, slot_a)
        sub_b = derive_agent_sub_seed(seed, DOMAIN_IDENTITY, slot_b)
        assert sub_a != sub_b

    def test_different_domains_different_sub_seeds(self) -> None:
        seed = mnemonic_to_seed(_TEST_MNEMONIC_12)
        slot = agent_id_to_slot("bot-alpha")
        DOMAIN_PAYMENTS = 1
        sub_id = derive_agent_sub_seed(seed, DOMAIN_IDENTITY, slot)
        sub_pay = derive_agent_sub_seed(seed, DOMAIN_PAYMENTS, slot)
        assert sub_id != sub_pay

    def test_sub_seed_differs_from_parent_seed(self) -> None:
        seed = mnemonic_to_seed(_TEST_MNEMONIC_12)
        slot = agent_id_to_slot("bot-alpha")
        sub = derive_agent_sub_seed(seed, DOMAIN_IDENTITY, slot)
        assert sub != seed


# ---------------------------------------------------------------------------
# run_recover — re-derive from mnemonic
# ---------------------------------------------------------------------------


class TestRunRecover:
    """muse auth recover must re-derive the exact same key from the mnemonic."""

    def _do_recover(
        self,
        monkeypatch: pytest.MonkeyPatch,
        tmp_path: pathlib.Path,
        mnemonic: str,
        extra_args: list[str] | None = None,
    ) -> object:
        fake_home = _patch_home(monkeypatch, tmp_path)
        args = ["auth", "recover", "--hub", _HUB] + (extra_args or [])
        return fake_home, runner.invoke(None, args, input=mnemonic)

    def test_recover_exits_zero(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        _, result = self._do_recover(monkeypatch, tmp_path, _TEST_MNEMONIC_12, ["--force"])
        assert result.exit_code == 0, result.output

    def test_recover_writes_pem(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        fake_home, result = self._do_recover(monkeypatch, tmp_path, _TEST_MNEMONIC_12, ["--force"])
        assert result.exit_code == 0
        pem = fake_home / ".muse" / "keys" / "localhost_10003.pem"
        assert pem.is_file()

    def test_recover_produces_same_fingerprint_as_keygen(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        """Key recovered from mnemonic must match the original keygen fingerprint."""
        import muse.core.bip39 as bip39_mod

        fixed_mnemonic = _TEST_MNEMONIC_12
        monkeypatch.setattr(bip39_mod, "generate_mnemonic", lambda **kw: fixed_mnemonic)

        # Keygen
        fake_home = _patch_home(monkeypatch, tmp_path)
        keygen_result = runner.invoke(None, ["auth", "keygen", "--hub", _HUB, "--json"])
        assert keygen_result.exit_code == 0, keygen_result.output
        keygen_fp = json.loads(keygen_result.output.splitlines()[0])["fingerprint"]

        # Recover into same tmpdir (--force to overwrite)
        recover_result = runner.invoke(
            None,
            ["auth", "recover", "--hub", _HUB, "--force", "--json"],
            input=fixed_mnemonic,
        )
        assert recover_result.exit_code == 0, recover_result.output
        recover_fp = json.loads(recover_result.output.splitlines()[0])["fingerprint"]

        assert keygen_fp == recover_fp, \
            f"Recovered fingerprint {recover_fp} != original {keygen_fp}"

    def test_recover_invalid_mnemonic_exits_nonzero(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        _, result = self._do_recover(monkeypatch, tmp_path, "not valid mnemonic words here ok", ["--force"])
        assert result.exit_code != 0

    def test_recover_pem_mode_600(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        fake_home, result = self._do_recover(monkeypatch, tmp_path, _TEST_MNEMONIC_12, ["--force"])
        assert result.exit_code == 0
        pem = fake_home / ".muse" / "keys" / "localhost_10003.pem"
        mode = pem.stat().st_mode & 0o777
        assert mode == 0o600

    def test_recover_json_has_fingerprint(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        _, result = self._do_recover(monkeypatch, tmp_path, _TEST_MNEMONIC_12, ["--force", "--json"])
        assert result.exit_code == 0
        payload = json.loads(result.output.splitlines()[0])
        assert "fingerprint" in payload
        assert len(payload["fingerprint"]) == 64


# ---------------------------------------------------------------------------
# Integration — full operator → agent → recover flow
# ---------------------------------------------------------------------------


class TestIntegrationFlow:
    """Full flow: human keygen → agent keygen → recover → fingerprints match."""

    def test_operator_then_agent_then_recover(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        import muse.core.bip39 as bip39_mod
        # Use a fixed mnemonic so we can recover without reading from the keychain.
        fixed_mnemonic = _TEST_MNEMONIC_12
        monkeypatch.setattr(bip39_mod, "generate_mnemonic", lambda **kw: fixed_mnemonic)

        _patch_home(monkeypatch, tmp_path)

        # 1. Operator keygen
        r1 = runner.invoke(None, ["auth", "keygen", "--hub", _HUB, "--json"])
        assert r1.exit_code == 0, r1.output
        op_payload = json.loads(r1.output.splitlines()[0])
        op_fp = op_payload["fingerprint"]

        # 2. Agent keygen derives from the operator's mnemonic in keychain / ephemeral store
        r2 = runner.invoke(
            None, ["auth", "keygen", "--hub", _HUB, "--agent-id", "worker-1", "--json"]
        )
        assert r2.exit_code == 0, r2.output
        agent_payload = json.loads(r2.output.splitlines()[0])
        agent_fp = agent_payload["fingerprint"]
        assert agent_fp != op_fp, "Agent fingerprint must differ from operator"

        # 3. Recover operator key via stdin pipe (--force since PEM already exists)
        r3 = runner.invoke(
            None,
            ["auth", "recover", "--hub", _HUB, "--force", "--json"],
            input=fixed_mnemonic,
        )
        assert r3.exit_code == 0, r3.output
        recovered_fp = json.loads(r3.output.splitlines()[0])["fingerprint"]
        assert recovered_fp == op_fp, \
            f"Recovered operator fp {recovered_fp!r} != original {op_fp!r}"

    def test_slot_stability_across_keygen_invocations(self) -> None:
        """agent_id_to_slot must return the same value before and after any keygen."""
        handle = "production-agent-42"
        slot_before = agent_id_to_slot(handle)
        # Simulate "after keygen" by just calling again — slot is a pure function
        slot_after = agent_id_to_slot(handle)
        assert slot_before == slot_after