"""Tests for memory zeroing in muse/cli/commands/agent.py key derivation.

Two functions in agent.py create DerivedKey objects without zeroing them:

1. _derive_agent_seed: calls mnemonic_to_seed() → bytearray, passes it to
   derive_agent_sub_seed(), but never zeros the intermediate seed bytearray.

2. _sub_seed_to_public: calls derive_identity_key(sub_seed) → DerivedKey,
   creates the Ed25519 key, but never calls dk.zero() on the DerivedKey.

Fix:
  - _derive_agent_seed: zero the seed bytearray after derive_agent_sub_seed().
  - _sub_seed_to_public: call dk.zero() after creating the Ed25519 key object.

Coverage
--------
I   _derive_agent_seed zeroes intermediate seed
    I1  seed bytearray is zeroed after derive_agent_sub_seed call

II  _sub_seed_to_public zeroes DerivedKey
    II1  DerivedKey.private_bytes is zeroed after _sub_seed_to_public returns
    II2  DerivedKey.chain_code is zeroed after _sub_seed_to_public returns
"""

from __future__ import annotations

from unittest.mock import patch

import pytest

from muse.core import hdkeys as _hdkeys
from muse.core.bip39 import mnemonic_to_seed
from muse.core.hdkeys import derive_agent_sub_seed, DOMAIN_IDENTITY
from muse.core.slip010 import DerivedKey
import muse.core.bip39 as _bip39_mod

_MNEMONIC = (
    "abandon abandon abandon abandon abandon abandon abandon abandon "
    "abandon abandon abandon about"
)
_SEED = mnemonic_to_seed(_MNEMONIC)
_SUB_SEED = derive_agent_sub_seed(_SEED, domain=DOMAIN_IDENTITY, agent_id=0)


# ---------------------------------------------------------------------------
# I  _derive_agent_seed zeroes intermediate seed
# ---------------------------------------------------------------------------

class TestDeriveAgentSeedZeroesIntermediateSeed:
    def test_I1_seed_zeroed_after_derive_agent_sub_seed(self) -> None:
        """I1: the bytearray seed from mnemonic_to_seed is zeroed after derive_agent_sub_seed."""
        from muse.cli.commands.agent import _derive_agent_seed

        seeds_captured: list[bytearray] = []
        original_to_seed = _bip39_mod.mnemonic_to_seed

        def tracking_to_seed(words, passphrase=""):
            result = original_to_seed(words, passphrase)
            seeds_captured.append(result)
            return result

        with patch.object(_bip39_mod, "mnemonic_to_seed", side_effect=tracking_to_seed):
            _derive_agent_seed(_MNEMONIC, account=0)

        assert seeds_captured, "mnemonic_to_seed was not called"
        seed = seeds_captured[0]
        assert seed == bytearray(64), (
            "intermediate seed bytearray must be zeroed after _derive_agent_seed returns"
        )


# ---------------------------------------------------------------------------
# II  _sub_seed_to_public zeroes DerivedKey
# ---------------------------------------------------------------------------

class TestSubSeedToPublicZeroesDerivedKey:
    def test_II1_private_bytes_zeroed(self) -> None:
        """II1: DerivedKey.private_bytes is zeroed after _sub_seed_to_public returns."""
        from muse.cli.commands.agent import _sub_seed_to_public

        captured: list[DerivedKey] = []
        original_derive = _hdkeys.derive_identity_key

        def capturing_derive(*args, **kwargs):
            dk = original_derive(*args, **kwargs)
            captured.append(dk)
            return dk

        with patch.object(_hdkeys, "derive_identity_key", side_effect=capturing_derive):
            _sub_seed_to_public(bytes(_SUB_SEED))

        assert captured, "derive_identity_key was not called"
        dk = captured[0]
        assert dk.private_bytes == bytearray(32), (
            "DerivedKey.private_bytes must be zeroed after _sub_seed_to_public"
        )

    def test_II2_chain_code_zeroed(self) -> None:
        """II2: DerivedKey.chain_code is zeroed after _sub_seed_to_public returns."""
        from muse.cli.commands.agent import _sub_seed_to_public

        captured: list[DerivedKey] = []
        original_derive = _hdkeys.derive_identity_key

        def capturing_derive(*args, **kwargs):
            dk = original_derive(*args, **kwargs)
            captured.append(dk)
            return dk

        with patch.object(_hdkeys, "derive_identity_key", side_effect=capturing_derive):
            _sub_seed_to_public(bytes(_SUB_SEED))

        assert captured, "derive_identity_key was not called"
        dk = captured[0]
        assert dk.chain_code == bytearray(32), (
            "DerivedKey.chain_code must be zeroed after _sub_seed_to_public"
        )