"""Phase 2A — HD mnemonic persistence tests.

Verifies that:
1. ``_load_all`` round-trips ``key_source``, ``mnemonic``, and ``hd_path``
   through identity.toml.
2. ``muse auth keygen --hd`` writes HD provenance to identity.toml immediately
   (mnemonic is NOT lost when the process exits).
3. ``muse auth register`` preserves HD fields from an existing identity entry
   rather than silently overwriting them with a plain ``IdentityEntry``.
4. Mnemonic is never emitted in the JSON stdout object (only on stderr).

Test categories covered
-----------------------
- unit          : _load_all / _dump_identity round-trip
- integration   : CLI round-trip via CliRunner
- e2e           : keygen → register field preservation
- stress        : 10 re-registrations, 20 successive save_identity writes
- data integrity: mnemonic survives repeated register calls + TOML escaping
- performance   : save+load under 100 ms; full keygen --hd under 3 s
- security      : mnemonic absent from JSON stdout object
- docstrings    : public API has docstrings (smoke)
"""

from __future__ import annotations

import json
import os
import pathlib
import tempfile
import time
import tomllib
from unittest.mock import patch

import pytest
from tests.cli_test_helper import CliRunner

from muse.core import keypair as kp_module

cli = None
runner = CliRunner()

HUB = "http://localhost:10003"
HOSTNAME = "localhost:10003"
FAKE_MNEMONIC = (
    "abandon abandon abandon abandon abandon abandon "
    "abandon abandon abandon abandon abandon about"
)
FAKE_HD_PATH = "m/1075233755'/0'/0'/0'/0'/0'"
FAKE_FINGERPRINT = "a" * 64   # used only in unit-test TOML fixtures (not real derivation)
FAKE_HANDLE = "gabriel"


# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------


def _patch_home(monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> pathlib.Path:
    """Redirect pathlib.Path.home() and module-level constants to a temp dir."""
    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")
    from muse.core import identity as id_module
    monkeypatch.setattr(id_module, "_IDENTITY_DIR", fake_home / ".muse")
    monkeypatch.setattr(id_module, "_IDENTITY_FILE", fake_home / ".muse" / "identity.toml")
    return fake_home


def _mock_hub(monkeypatch: pytest.MonkeyPatch, handle: str = FAKE_HANDLE) -> None:
    """Patch _json_post_raw to simulate a successful hub challenge-response."""
    import muse.cli.commands.auth as auth_mod
    challenge = {"challengeToken": "deadbeef" * 8, "isNewKey": True, "algorithm": "ed25519"}
    verify = {"handle": handle, "identityId": "id-123", "isNewIdentity": False, "authMethod": "ed25519"}
    monkeypatch.setattr(
        auth_mod,
        "_json_post_raw",
        lambda base, path, payload: challenge if "challenge" in path else verify,
    )


def _mock_bip39(monkeypatch: pytest.MonkeyPatch) -> None:
    """Patch only mnemonic *generation* to avoid OS entropy.

    ``mnemonic_to_seed`` and ``generate_hd_keypair`` run for real so that:
    - A valid PEM key is written to disk (run_register can sign with it).
    - The fingerprint stored in identity.toml is the genuine derived value.
    - SLIP-0010 derivation is exercised, not bypassed.
    """
    import muse.core.bip39 as bip39_mod
    monkeypatch.setattr(bip39_mod, "generate_mnemonic", lambda **kw: FAKE_MNEMONIC)


def _identity_file(fake_home: pathlib.Path) -> pathlib.Path:
    return fake_home / ".muse" / "identity.toml"


def _read_identity_toml(fake_home: pathlib.Path) -> dict:
    ifile = _identity_file(fake_home)
    with ifile.open("rb") as fh:
        return tomllib.load(fh)


# ---------------------------------------------------------------------------
# 1. Unit — _load_all round-trips HD fields
# ---------------------------------------------------------------------------


class TestLoadAllHdFields:
    """_load_all must parse key_source, mnemonic, hd_path from TOML."""

    def _write(self, path: pathlib.Path, text: str) -> None:
        path.write_text(text, encoding="utf-8")

    def test_loads_key_source(self, tmp_path):
        p = tmp_path / "identity.toml"
        self._write(p, f'["{HOSTNAME}"]\ntype="human"\nhandle="{FAKE_HANDLE}"\n'
                       f'key_path="/tmp/k.pem"\nalgorithm="ed25519"\n'
                       f'fingerprint="{FAKE_FINGERPRINT}"\nkey_source="hd"\n')
        from muse.core.identity import _load_all
        assert _load_all(p)[HOSTNAME]["key_source"] == "hd"

    def test_loads_mnemonic(self, tmp_path):
        p = tmp_path / "identity.toml"
        self._write(p, f'["{HOSTNAME}"]\ntype="human"\nhandle="{FAKE_HANDLE}"\n'
                       f'key_path="/tmp/k.pem"\nalgorithm="ed25519"\n'
                       f'fingerprint="{FAKE_FINGERPRINT}"\nmnemonic="{FAKE_MNEMONIC}"\n')
        from muse.core.identity import _load_all
        assert _load_all(p)[HOSTNAME]["mnemonic"] == FAKE_MNEMONIC

    def test_loads_hd_path(self, tmp_path):
        p = tmp_path / "identity.toml"
        self._write(p, f'["{HOSTNAME}"]\ntype="human"\nhandle="{FAKE_HANDLE}"\n'
                       f'key_path="/tmp/k.pem"\nalgorithm="ed25519"\n'
                       f'fingerprint="{FAKE_FINGERPRINT}"\nhd_path="{FAKE_HD_PATH}"\n')
        from muse.core.identity import _load_all
        assert _load_all(p)[HOSTNAME]["hd_path"] == FAKE_HD_PATH

    def test_loads_all_hd_fields_together(self, tmp_path):
        p = tmp_path / "identity.toml"
        self._write(p, f'["{HOSTNAME}"]\ntype="human"\nhandle="{FAKE_HANDLE}"\n'
                       f'key_path="/tmp/k.pem"\nalgorithm="ed25519"\n'
                       f'fingerprint="{FAKE_FINGERPRINT}"\n'
                       f'key_source="hd"\nmnemonic="{FAKE_MNEMONIC}"\nhd_path="{FAKE_HD_PATH}"\n')
        from muse.core.identity import _load_all
        entry = _load_all(p)[HOSTNAME]
        assert entry["key_source"] == "hd"
        assert entry["mnemonic"] == FAKE_MNEMONIC
        assert entry["hd_path"] == FAKE_HD_PATH

    def test_jbok_entry_has_no_hd_fields(self, tmp_path):
        p = tmp_path / "identity.toml"
        self._write(p, f'["{HOSTNAME}"]\ntype="human"\nhandle="{FAKE_HANDLE}"\n'
                       f'key_path="/tmp/k.pem"\nalgorithm="ed25519"\n'
                       f'fingerprint="{FAKE_FINGERPRINT}"\n')
        from muse.core.identity import _load_all
        entry = _load_all(p)[HOSTNAME]
        assert "key_source" not in entry
        assert "mnemonic" not in entry
        assert "hd_path" not in entry

    def test_round_trip_hd_fields(self, tmp_path):
        """_dump_identity → write → _load_all preserves HD fields exactly."""
        from muse.core.identity import _dump_identity, _load_all
        identities = {HOSTNAME: {
            "type": "human", "handle": FAKE_HANDLE,
            "key_path": "/tmp/k.pem", "algorithm": "ed25519",
            "fingerprint": FAKE_FINGERPRINT,
            "key_source": "hd", "mnemonic": FAKE_MNEMONIC, "hd_path": FAKE_HD_PATH,
        }}
        p = tmp_path / "identity.toml"
        p.write_text(_dump_identity(identities), encoding="utf-8")
        entry = _load_all(p)[HOSTNAME]
        assert entry["key_source"] == "hd"
        assert entry["mnemonic"] == FAKE_MNEMONIC
        assert entry["hd_path"] == FAKE_HD_PATH


# ---------------------------------------------------------------------------
# 2. Integration — run_keygen --hd writes identity.toml
# ---------------------------------------------------------------------------


class TestKeygenHdWritesIdentity:
    """run_keygen --hd must persist HD fields to identity.toml."""

    def _run(self, monkeypatch, tmp_path, extra_args=None):
        fake_home = _patch_home(monkeypatch, tmp_path)
        _mock_bip39(monkeypatch)
        args = ["auth", "keygen", "--hub", HUB, "--hd"] + (extra_args or [])
        result = runner.invoke(cli, args, catch_exceptions=False)
        return result, fake_home

    def test_identity_toml_created(self, monkeypatch, tmp_path):
        result, fake_home = self._run(monkeypatch, tmp_path)
        assert result.exit_code == 0, result.output
        assert _identity_file(fake_home).exists()

    def test_key_source_hd_written(self, monkeypatch, tmp_path):
        _, fake_home = self._run(monkeypatch, tmp_path)
        data = _read_identity_toml(fake_home)
        assert data[HOSTNAME]["key_source"] == "hd"

    def test_mnemonic_written(self, monkeypatch, tmp_path):
        _, fake_home = self._run(monkeypatch, tmp_path)
        data = _read_identity_toml(fake_home)
        assert data[HOSTNAME]["mnemonic"] == FAKE_MNEMONIC

    def test_hd_path_written(self, monkeypatch, tmp_path):
        _, fake_home = self._run(monkeypatch, tmp_path)
        data = _read_identity_toml(fake_home)
        assert data[HOSTNAME]["hd_path"].startswith("m/")

    def test_json_output_has_no_mnemonic_key(self, monkeypatch, tmp_path):
        """JSON stdout object must not contain a 'mnemonic' key."""
        result, _ = self._run(monkeypatch, tmp_path, extra_args=["--json"])
        assert result.exit_code == 0
        # Find the JSON line in combined output (stdout is first, before stderr)
        json_line = next(
            (line for line in result.output.splitlines() if line.startswith("{")), None
        )
        assert json_line is not None, "No JSON in output"
        out = json.loads(json_line)
        assert "mnemonic" not in out

    def test_mnemonic_absent_from_json_object(self, monkeypatch, tmp_path):
        """Even when mnemonic appears in stderr, JSON dict must not carry it."""
        result, _ = self._run(monkeypatch, tmp_path, extra_args=["--json"])
        json_line = next(
            (line for line in result.output.splitlines() if line.startswith("{")), None
        )
        out = json.loads(json_line)
        assert FAKE_MNEMONIC not in json.dumps(out)

    def test_keygen_hd_force_overwrites_identity(self, monkeypatch, tmp_path):
        """--force on an existing HD key overwrites identity.toml entry."""
        self._run(monkeypatch, tmp_path)
        result, fake_home = self._run(monkeypatch, tmp_path, extra_args=["--force"])
        assert result.exit_code == 0, result.output
        data = _read_identity_toml(fake_home)
        assert data[HOSTNAME]["key_source"] == "hd"


# ---------------------------------------------------------------------------
# 3. E2E — run_register preserves HD fields
# ---------------------------------------------------------------------------


class TestRegisterPreservesHdFields:
    """run_register must carry forward key_source, mnemonic, hd_path."""

    def _setup_hd_keygen(self, monkeypatch, tmp_path) -> pathlib.Path:
        """Run keygen --hd so identity.toml gets HD fields, return fake_home.

        Only ``generate_mnemonic`` is mocked (to avoid OS entropy and
        non-determinism); ``mnemonic_to_seed`` and ``generate_hd_keypair``
        run for real so that a valid PEM is written to disk and
        run_register can actually sign the challenge.
        """
        fake_home = _patch_home(monkeypatch, tmp_path)
        _mock_bip39(monkeypatch)
        args = ["auth", "keygen", "--hub", HUB, "--hd"]
        result = runner.invoke(cli, args, catch_exceptions=False)
        assert result.exit_code == 0, result.output
        return fake_home

    def _run_register(self, monkeypatch, fake_home) -> "object":
        _mock_hub(monkeypatch)
        args = ["auth", "register", "--hub", HUB, "--handle", FAKE_HANDLE]
        result = runner.invoke(cli, args, catch_exceptions=False)
        return result

    def test_hd_fields_preserved_after_register(self, monkeypatch, tmp_path):
        fake_home = self._setup_hd_keygen(monkeypatch, tmp_path)
        result = self._run_register(monkeypatch, fake_home)
        assert result.exit_code == 0, result.output

        data = _read_identity_toml(fake_home)
        entry = data[HOSTNAME]
        assert entry.get("key_source") == "hd", "key_source lost after register"
        assert entry.get("mnemonic") == FAKE_MNEMONIC, "mnemonic lost after register"
        assert entry.get("hd_path", "").startswith("m/"), "hd_path lost after register"

    def test_handle_updated_after_register(self, monkeypatch, tmp_path):
        """Provisional empty handle must be replaced by server-assigned handle."""
        fake_home = self._setup_hd_keygen(monkeypatch, tmp_path)
        # Confirm provisional entry has empty handle before register
        pre_data = _read_identity_toml(fake_home)
        assert pre_data[HOSTNAME].get("handle", "") == ""

        self._run_register(monkeypatch, fake_home)
        post_data = _read_identity_toml(fake_home)
        assert post_data[HOSTNAME].get("handle") == FAKE_HANDLE

    def test_register_second_time_still_preserves_mnemonic(self, monkeypatch, tmp_path):
        """Re-authentication preserves the mnemonic through multiple registrations."""
        fake_home = self._setup_hd_keygen(monkeypatch, tmp_path)
        self._run_register(monkeypatch, fake_home)
        # Second registration (e.g. key rotation to same hub)
        result2 = self._run_register(monkeypatch, fake_home)
        assert result2.exit_code == 0, result2.output

        data = _read_identity_toml(fake_home)
        assert data[HOSTNAME].get("mnemonic") == FAKE_MNEMONIC

    def test_jbok_register_no_hd_fields_added(self, monkeypatch, tmp_path):
        """Registering with a JBOK key must not add HD fields."""
        fake_home = _patch_home(monkeypatch, tmp_path)
        # Generate a JBOK key (not HD)
        kp_module.generate_keypair(HOSTNAME)
        _mock_hub(monkeypatch)
        result = runner.invoke(
            cli,
            ["auth", "register", "--hub", HUB, "--handle", FAKE_HANDLE],
            catch_exceptions=False,
        )
        assert result.exit_code == 0, result.output

        if _identity_file(fake_home).exists():
            data = _read_identity_toml(fake_home)
            entry = data.get(HOSTNAME, {})
            assert "key_source" not in entry
            assert "mnemonic" not in entry


# ---------------------------------------------------------------------------
# 4. Data integrity — mnemonic survives TOML escaping edge cases
# ---------------------------------------------------------------------------


class TestMnemonicTomlEscaping:
    def test_mnemonic_with_quotes_round_trips(self, tmp_path):
        """A mnemonic containing TOML-special chars survives _dump → _load."""
        from muse.core.identity import _dump_identity, _load_all
        weird = 'word1 word2 "quoted" word3 back\\slash word4'
        entry = {
            "type": "human", "handle": FAKE_HANDLE,
            "key_path": "/tmp/k.pem", "algorithm": "ed25519",
            "fingerprint": FAKE_FINGERPRINT,
            "key_source": "hd", "mnemonic": weird, "hd_path": FAKE_HD_PATH,
        }
        p = tmp_path / "identity.toml"
        p.write_text(_dump_identity({HOSTNAME: entry}), encoding="utf-8")
        assert _load_all(p)[HOSTNAME]["mnemonic"] == weird

    def test_hd_path_prime_and_slash_preserved(self, tmp_path):
        """HD path with primes and slashes round-trips without corruption."""
        from muse.core.identity import _dump_identity, _load_all
        entry = {
            "type": "human", "handle": FAKE_HANDLE,
            "key_path": "/tmp/k.pem", "algorithm": "ed25519",
            "fingerprint": FAKE_FINGERPRINT,
            "key_source": "hd", "mnemonic": FAKE_MNEMONIC, "hd_path": FAKE_HD_PATH,
        }
        p = tmp_path / "identity.toml"
        p.write_text(_dump_identity({HOSTNAME: entry}), encoding="utf-8")
        assert _load_all(p)[HOSTNAME]["hd_path"] == FAKE_HD_PATH


# ---------------------------------------------------------------------------
# 5. Security — mnemonic never in JSON stdout object
# ---------------------------------------------------------------------------


class TestMnemonicNeverInJsonObject:
    """Mnemonic must not appear in any JSON stdout object."""

    def test_keygen_hd_json_no_mnemonic_key(self, monkeypatch, tmp_path):
        _patch_home(monkeypatch, tmp_path)
        _mock_bip39(monkeypatch)
        result = runner.invoke(
            cli,
            ["auth", "keygen", "--hub", HUB, "--hd", "--json"],
            catch_exceptions=False,
        )
        assert result.exit_code == 0
        json_line = next(
            (l for l in result.output.splitlines() if l.startswith("{")), None
        )
        assert json_line is not None
        obj = json.loads(json_line)
        assert "mnemonic" not in obj

    def test_keygen_hd_json_mnemonic_word_count_present(self, monkeypatch, tmp_path):
        """JSON should have mnemonic_word_count (count, not content)."""
        _patch_home(monkeypatch, tmp_path)
        _mock_bip39(monkeypatch)
        result = runner.invoke(
            cli,
            ["auth", "keygen", "--hub", HUB, "--hd", "--json"],
            catch_exceptions=False,
        )
        json_line = next(
            (l for l in result.output.splitlines() if l.startswith("{")), None
        )
        obj = json.loads(json_line)
        assert "mnemonic_word_count" in obj
        assert obj["mnemonic_word_count"] == 12  # 128-bit → 12 words


# ---------------------------------------------------------------------------
# 6. Docstring smoke tests
# ---------------------------------------------------------------------------


class TestDocstrings:
    def test_load_all_has_docstring(self):
        from muse.core.identity import _load_all
        assert _load_all.__doc__

    def test_save_identity_has_docstring(self):
        from muse.core.identity import save_identity
        assert save_identity.__doc__

    def test_load_identity_has_docstring(self):
        from muse.core.identity import load_identity
        assert load_identity.__doc__

    def test_generate_hd_keypair_has_docstring(self):
        from muse.core.keypair import generate_hd_keypair
        assert generate_hd_keypair.__doc__


# ---------------------------------------------------------------------------
# Performance
# ---------------------------------------------------------------------------


class TestPersistencePerformance:
    """Identity TOML read/write must add negligible latency."""

    def test_save_and_load_identity_under_100ms(self, tmp_path):
        """save_identity + load_identity must complete in under 100 ms."""
        from muse.core import identity as id_module
        import pytest
        identity_file = tmp_path / "identity.toml"
        monkeypatch = pytest.MonkeyPatch()
        monkeypatch.setattr(id_module, "_IDENTITY_DIR", tmp_path)
        monkeypatch.setattr(id_module, "_IDENTITY_FILE", identity_file)
        entry = {
            "type": "human", "handle": FAKE_HANDLE,
            "key_path": "/tmp/k.pem", "algorithm": "ed25519",
            "fingerprint": FAKE_FINGERPRINT,
            "key_source": "hd", "mnemonic": FAKE_MNEMONIC, "hd_path": FAKE_HD_PATH,
        }
        start = time.monotonic()
        id_module.save_identity(HUB, entry)
        loaded = id_module.load_identity(HUB)
        elapsed = time.monotonic() - start
        monkeypatch.undo()
        assert loaded is not None
        assert elapsed < 0.1, f"save+load took {elapsed*1000:.1f}ms"

    def test_keygen_hd_then_load_identity_under_3s(self, monkeypatch, tmp_path):
        """The full keygen --hd path including SLIP-0010 must complete in under 3 s."""
        _patch_home(monkeypatch, tmp_path)
        _mock_bip39(monkeypatch)
        start = time.monotonic()
        result = runner.invoke(
            cli,
            ["auth", "keygen", "--hub", HUB, "--hd"],
            catch_exceptions=False,
        )
        elapsed = time.monotonic() - start
        assert result.exit_code == 0, result.output
        assert elapsed < 3.0, f"keygen --hd took {elapsed:.2f}s"


# ---------------------------------------------------------------------------
# Stress
# ---------------------------------------------------------------------------


class TestPersistenceStress:
    """HD field persistence must hold under repeated writes and re-loads."""

    def test_10_successive_registers_preserve_mnemonic(self, monkeypatch, tmp_path):
        """Mnemonic must survive 10 consecutive register calls without corruption."""
        fake_home = _patch_home(monkeypatch, tmp_path)
        _mock_bip39(monkeypatch)
        # Generate HD key once
        runner.invoke(cli, ["auth", "keygen", "--hub", HUB, "--hd"],
                      catch_exceptions=False)
        # Simulate 10 re-registrations
        _mock_hub(monkeypatch)
        for i in range(10):
            result = runner.invoke(
                cli,
                ["auth", "register", "--hub", HUB, "--handle", FAKE_HANDLE],
                catch_exceptions=False,
            )
            assert result.exit_code == 0, f"iteration {i}: {result.output}"
            data = _read_identity_toml(fake_home)
            mnemonic_stored = data[HOSTNAME].get("mnemonic", "")
            assert mnemonic_stored == FAKE_MNEMONIC, \
                f"Mnemonic corrupted after {i+1} register calls"

    def test_concurrent_identity_writes_do_not_corrupt(self, monkeypatch, tmp_path):
        """Multiple save_identity calls in succession must not corrupt the file."""
        from muse.core import identity as id_module
        identity_file = tmp_path / "identity.toml"
        monkeypatch.setattr(id_module, "_IDENTITY_DIR", tmp_path)
        monkeypatch.setattr(id_module, "_IDENTITY_FILE", identity_file)

        base_entry = {
            "type": "human", "handle": FAKE_HANDLE,
            "key_path": "/tmp/k.pem", "algorithm": "ed25519",
            "fingerprint": FAKE_FINGERPRINT,
            "key_source": "hd", "mnemonic": FAKE_MNEMONIC, "hd_path": FAKE_HD_PATH,
        }
        for i in range(20):
            entry = {**base_entry, "handle": f"user_{i}"}
            id_module.save_identity(HUB, entry)

        loaded = id_module.load_identity(HUB)
        assert loaded is not None
        assert loaded.get("key_source") == "hd"
        assert loaded.get("mnemonic") == FAKE_MNEMONIC