"""Tests for b2alpha.crypto – key generation, signing, verification."""

from __future__ import annotations

import base64
import hashlib
import json

import pytest
from nacl.signing import SigningKey, VerifyKey


# ---------------------------------------------------------------------------
# Helpers (mirror the module's b64url helpers to verify independently)
# ---------------------------------------------------------------------------

def _b64url_decode(s: str) -> bytes:
    padding = 4 - len(s) % 4
    if padding != 4:
        s += "=" * padding
    return base64.urlsafe_b64decode(s)


# ---------------------------------------------------------------------------
# generate_keypair
# ---------------------------------------------------------------------------

class TestGenerateKeypair:
    def test_returns_did_pubkey_path(self, tmp_path, monkeypatch):
        import b2alpha.config as cfg
        monkeypatch.setattr(cfg, "KEYS_DIR", tmp_path / "keys")

        from b2alpha.crypto import generate_keypair
        did, pub_key, key_path = generate_keypair("test_agent")

        assert did.startswith("did:b2a:z")
        assert len(pub_key) > 0
        assert key_path.endswith("test_agent.key")

    def test_key_file_is_valid_ed25519(self, tmp_path, monkeypatch):
        import b2alpha.config as cfg
        monkeypatch.setattr(cfg, "KEYS_DIR", tmp_path / "keys")

        from b2alpha.crypto import generate_keypair
        _, pub_b64, key_path = generate_keypair("ed_check")

        from pathlib import Path
        raw = Path(key_path).read_bytes()
        assert len(raw) == 32  # Ed25519 seed size
        sk = SigningKey(raw)
        assert bytes(sk.verify_key) == _b64url_decode(pub_b64)

    def test_did_derived_from_public_key(self, tmp_path, monkeypatch):
        import b2alpha.config as cfg
        monkeypatch.setattr(cfg, "KEYS_DIR", tmp_path / "keys")

        from b2alpha.crypto import generate_keypair, _b64url_encode
        did, pub_b64, _ = generate_keypair("did_check")

        pk_bytes = _b64url_decode(pub_b64)
        expected_did = "did:b2a:z" + _b64url_encode(pk_bytes)[:22]
        assert did == expected_did

    def test_two_keypairs_differ(self, tmp_path, monkeypatch):
        import b2alpha.config as cfg
        monkeypatch.setattr(cfg, "KEYS_DIR", tmp_path / "keys")

        from b2alpha.crypto import generate_keypair
        did1, _, _ = generate_keypair("a1")
        did2, _, _ = generate_keypair("a2")
        assert did1 != did2

    def test_key_file_permissions(self, tmp_path, monkeypatch):
        import os
        import b2alpha.config as cfg
        monkeypatch.setattr(cfg, "KEYS_DIR", tmp_path / "keys")

        from b2alpha.crypto import generate_keypair
        _, _, key_path = generate_keypair("perm_check")

        from pathlib import Path
        mode = Path(key_path).stat().st_mode & 0o777
        assert mode == 0o600


# ---------------------------------------------------------------------------
# load_signing_key
# ---------------------------------------------------------------------------

class TestLoadSigningKey:
    def test_loads_generated_key(self, tmp_path, monkeypatch):
        import b2alpha.config as cfg
        monkeypatch.setattr(cfg, "KEYS_DIR", tmp_path / "keys")

        from b2alpha.crypto import generate_keypair, load_signing_key
        generate_keypair("roundtrip")
        sk = load_signing_key("roundtrip")
        assert isinstance(sk, SigningKey)

    def test_missing_key_raises(self, tmp_path, monkeypatch):
        import b2alpha.config as cfg
        monkeypatch.setattr(cfg, "KEYS_DIR", tmp_path / "keys")

        from b2alpha.crypto import load_signing_key
        with pytest.raises(FileNotFoundError, match="No key found"):
            load_signing_key("nonexistent")


# ---------------------------------------------------------------------------
# sign_message
# ---------------------------------------------------------------------------

class TestSignMessage:
    @pytest.fixture(autouse=True)
    def _setup_keys(self, tmp_path, monkeypatch):
        import b2alpha.config as cfg
        monkeypatch.setattr(cfg, "KEYS_DIR", tmp_path / "keys")
        from b2alpha.crypto import generate_keypair
        # Generate as "default" (used by sign_message's default key_name)
        self.did, self.pub_b64, _ = generate_keypair("default")

    def test_returns_required_fields(self):
        from b2alpha.crypto import sign_message
        env = sign_message(self.did, "did:b2a:zRecipient0000", b"hello")
        assert "message_id" in env
        assert "payload" in env
        assert "signature" in env
        assert env["protocol_version"] == "0.1"

    def test_custom_message_id(self):
        from b2alpha.crypto import sign_message
        env = sign_message(self.did, "did:b2a:zRecip", b"data", message_id="custom-123")
        assert env["message_id"] == "custom-123"

    def test_payload_roundtrips(self):
        from b2alpha.crypto import sign_message
        original = json.dumps({"text": "hello world"}).encode()
        env = sign_message(self.did, "did:b2a:zRecip", original)
        decoded = _b64url_decode(env["payload"])
        assert decoded == original

    def test_signature_verifies(self):
        from b2alpha.crypto import sign_message
        recipient = "did:b2a:zRecipientXXX"
        payload = b"verify me"
        env = sign_message(self.did, recipient, payload, key_name="default")

        # Reconstruct the signed data the same way the module does
        to_sign = (
            env["message_id"].encode()
            + self.did.encode()
            + recipient.encode()
            + payload
        )
        digest = hashlib.sha256(to_sign).digest()

        sig_bytes = _b64url_decode(env["signature"])
        pk_bytes = _b64url_decode(self.pub_b64)
        vk = VerifyKey(pk_bytes)
        # Raises nacl.exceptions.BadSignatureError if invalid
        vk.verify(digest, sig_bytes)

    def test_tampered_payload_fails_verification(self):
        from nacl.exceptions import BadSignatureError
        from b2alpha.crypto import sign_message

        recipient = "did:b2a:zRecipientXXX"
        env = sign_message(self.did, recipient, b"original", key_name="default")

        # Tamper: verify with different payload
        to_sign = (
            env["message_id"].encode()
            + self.did.encode()
            + recipient.encode()
            + b"TAMPERED"
        )
        digest = hashlib.sha256(to_sign).digest()
        sig_bytes = _b64url_decode(env["signature"])
        pk_bytes = _b64url_decode(self.pub_b64)
        vk = VerifyKey(pk_bytes)

        with pytest.raises(BadSignatureError):
            vk.verify(digest, sig_bytes)


# ---------------------------------------------------------------------------
# b64url encode/decode roundtrip
# ---------------------------------------------------------------------------

class TestB64Url:
    @pytest.mark.parametrize("data", [b"", b"a", b"ab", b"abc", b"\x00\xff" * 20])
    def test_roundtrip(self, data):
        from b2alpha.crypto import _b64url_encode, _b64url_decode
        assert _b64url_decode(_b64url_encode(data)) == data
