"""Tests for cryptographic message signing."""

import json
import time
import pytest

from crewai_mcps.secure_tool import KeyManager, NonceTracker, SecureMCPTool


class TestKeyManager:
    """Test key generation and signing operations."""

    def test_auto_generates_keys(self):
        km = KeyManager()
        assert km.private_key is not None
        assert km.public_key is not None

    def test_public_key_pem_export(self):
        km = KeyManager()
        pem = km.public_key_pem()
        assert pem.startswith("-----BEGIN PUBLIC KEY-----")
        assert pem.strip().endswith("-----END PUBLIC KEY-----")

    def test_sign_produces_base64(self):
        km = KeyManager()
        sig = km.sign(b"test data")
        assert isinstance(sig, str)
        assert len(sig) > 0
        # Should be valid base64
        import base64
        decoded = base64.b64decode(sig)
        assert len(decoded) > 0

    def test_sign_and_verify_roundtrip(self):
        km = KeyManager()
        data = b"hello world"
        sig = km.sign(data)
        assert km.verify(data, sig) is True

    def test_verify_fails_with_wrong_data(self):
        km = KeyManager()
        sig = km.sign(b"correct data")
        assert km.verify(b"wrong data", sig) is False

    def test_verify_fails_with_wrong_key(self):
        km1 = KeyManager()
        km2 = KeyManager()
        sig = km1.sign(b"test")
        assert km2.verify(b"test", sig) is False

    def test_verify_fails_with_invalid_signature(self):
        km = KeyManager()
        assert km.verify(b"test", "not-a-valid-signature") is False

    def test_different_data_different_signatures(self):
        km = KeyManager()
        sig1 = km.sign(b"data one")
        sig2 = km.sign(b"data two")
        assert sig1 != sig2

    def test_custom_key_injection(self):
        from cryptography.hazmat.primitives.asymmetric import ec
        custom_key = ec.generate_private_key(ec.SECP256R1())
        km = KeyManager(private_key=custom_key)
        sig = km.sign(b"test")
        assert km.verify(b"test", sig) is True


class TestNonceTracker:
    """Test nonce generation and replay detection."""

    def test_generate_unique_nonces(self):
        tracker = NonceTracker()
        nonces = {tracker.generate() for _ in range(100)}
        assert len(nonces) == 100

    def test_nonce_is_hex_string(self):
        tracker = NonceTracker()
        nonce = tracker.generate()
        assert isinstance(nonce, str)
        assert len(nonce) == 32  # 16 bytes = 32 hex chars
        int(nonce, 16)  # Should not raise

    def test_replay_detection(self):
        tracker = NonceTracker()
        nonce = tracker.generate()
        # First check: nonce is known (was generated)
        assert tracker.is_known(nonce) is True
        # Trying to consume an already-used nonce should fail
        assert tracker.check_and_consume(nonce) is False

    def test_fresh_nonce_accepted(self):
        tracker = NonceTracker()
        # Manually check a never-seen nonce
        assert tracker.check_and_consume("aaaa" * 8) is True
        # Second use should be rejected
        assert tracker.check_and_consume("aaaa" * 8) is False

    def test_expired_nonces_cleaned(self):
        tracker = NonceTracker(max_age_seconds=0)
        nonce = tracker.generate()
        time.sleep(0.01)
        tracker._cleanup()
        # After expiry, a fresh check_and_consume should succeed
        # because the nonce was cleaned out
        assert tracker.check_and_consume(nonce) is True


class TestSecureMCPToolSigning:
    """Test the full signing workflow."""

    def test_sign_call_produces_envelope(self):
        secure = SecureMCPTool()
        envelope = secure.sign_call("test_tool", {"query": "hello"})
        assert "nonce" in envelope
        assert "timestamp" in envelope
        assert "signature" in envelope
        assert "tool_name" in envelope
        assert envelope["tool_name"] == "test_tool"

    def test_sign_call_unique_nonces(self):
        secure = SecureMCPTool()
        nonces = set()
        for _ in range(50):
            envelope = secure.sign_call("tool", {"x": 1})
            nonces.add(envelope["nonce"])
        assert len(nonces) == 50

    def test_signature_is_verifiable(self):
        secure = SecureMCPTool()
        envelope = secure.sign_call("my_tool", {"arg": "value"})

        # Reconstruct the payload and verify the signature
        payload = {
            "tool_name": "my_tool",
            "arguments": {"arg": "value"},
            "nonce": envelope["nonce"],
            "timestamp": envelope["timestamp"],
        }
        canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"))
        assert secure.keys.verify(
            canonical.encode("utf-8"), envelope["signature"]
        ) is True

    def test_signature_fails_with_tampered_args(self):
        secure = SecureMCPTool()
        envelope = secure.sign_call("tool", {"safe": True})

        # Tamper with the arguments
        payload = {
            "tool_name": "tool",
            "arguments": {"safe": False},  # Tampered!
            "nonce": envelope["nonce"],
            "timestamp": envelope["timestamp"],
        }
        canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"))
        assert secure.keys.verify(
            canonical.encode("utf-8"), envelope["signature"]
        ) is False

    def test_verify_response(self):
        secure = SecureMCPTool()
        envelope = secure.sign_call("tool", {"x": 1})
        result = secure.verify_response("tool", "response data", envelope)
        assert result["verified"] is True
        assert result["tool_name"] == "tool"
        assert result["nonce"] == envelope["nonce"]
