"""Tests for the 4-layer security module."""
import json
import os
import time
import tempfile
import pytest

from greymatter_core.security.tokens import (
    generate_node_token,
    verify_node_token,
    TokenStore,
)
from greymatter_core.security.encryption import (
    encrypt_sync_payload,
    decrypt_sync_payload,
    get_sqlcipher_key,
    _derive_key,
    MLKEM_AVAILABLE,
)

# Conditional imports for ML-KEM tests
if MLKEM_AVAILABLE:
    from greymatter_core.security.encryption import (
        generate_mlkem_keypair,
        save_mlkem_keys,
        load_mlkem_keys,
    )


# =========================================================================
# Layer 1: Token Tests
# =========================================================================

class TestTokenGeneration:
    SECRET = b"test-cluster-secret-32bytes-long!"

    def test_generate_returns_four_parts(self):
        token = generate_node_token("node-1", secret=self.SECRET)
        parts = token.split(".")
        assert len(parts) == 4
        assert parts[0] == "node-1"

    def test_verify_valid_token(self):
        token = generate_node_token("node-1", secret=self.SECRET)
        valid, node_id = verify_node_token(token, secret=self.SECRET)
        assert valid is True
        assert node_id == "node-1"

    def test_verify_wrong_secret(self):
        token = generate_node_token("node-1", secret=self.SECRET)
        valid, node_id = verify_node_token(token, secret=b"wrong-secret")
        assert valid is False
        assert node_id == ""

    def test_verify_tampered_token(self):
        token = generate_node_token("node-1", secret=self.SECRET)
        # Tamper with the signature (last part)
        parts = token.split(".")
        parts[3] = "0" * 64
        tampered = ".".join(parts)
        valid, _ = verify_node_token(tampered, secret=self.SECRET)
        assert valid is False

    def test_verify_expired_token(self):
        token = generate_node_token("node-1", secret=self.SECRET)
        # Token with 0 max_age should immediately expire
        valid, _ = verify_node_token(token, secret=self.SECRET, max_age=0)
        assert valid is False

    def test_verify_malformed_token(self):
        assert verify_node_token("garbage", secret=self.SECRET) == (False, "")
        assert verify_node_token("a.b", secret=self.SECRET) == (False, "")
        assert verify_node_token("", secret=self.SECRET) == (False, "")

    def test_different_nodes_different_tokens(self):
        t1 = generate_node_token("node-1", secret=self.SECRET)
        t2 = generate_node_token("node-2", secret=self.SECRET)
        assert t1 != t2

    def test_constant_time_comparison(self):
        """Verify we use hmac.compare_digest (timing-safe)."""
        token = generate_node_token("node-1", secret=self.SECRET)
        # This test just ensures the function works correctly;
        # actual timing analysis requires specialized tools.
        valid, _ = verify_node_token(token, secret=self.SECRET)
        assert valid is True


class TestTokenStore:
    def test_issue_and_verify(self, tmp_path):
        store = TokenStore(str(tmp_path / "tokens.db"))
        token = store.issue("node-1")
        valid, node_id = store.verify(token)
        assert valid is True
        assert node_id == "node-1"
        store.close()

    def test_rotation_preserves_old_token(self, tmp_path):
        store = TokenStore(str(tmp_path / "tokens.db"))
        token1 = store.issue("node-1")
        token2 = store.issue("node-1")  # Rotate
        assert token1 != token2
        # Both should verify (token1 is still valid by HMAC)
        assert store.verify(token1)[0] is True
        assert store.verify(token2)[0] is True
        store.close()

    def test_revoke(self, tmp_path):
        store = TokenStore(str(tmp_path / "tokens.db"))
        store.issue("node-1")
        store.revoke("node-1")
        nodes = store.list_nodes()
        assert len(nodes) == 0
        store.close()

    def test_list_nodes(self, tmp_path):
        store = TokenStore(str(tmp_path / "tokens.db"))
        store.issue("node-1")
        store.issue("node-2")
        nodes = store.list_nodes()
        assert len(nodes) == 2
        node_ids = {n["node_id"] for n in nodes}
        assert node_ids == {"node-1", "node-2"}
        store.close()


# =========================================================================
# Layer 3: Encryption at Rest Tests
# =========================================================================

class TestEncryptionAtRest:
    def test_sqlcipher_key_unique_per_node(self):
        secret = b"test-secret"
        key1 = get_sqlcipher_key("node-1", secret=secret)
        key2 = get_sqlcipher_key("node-2", secret=secret)
        assert key1 != key2
        assert len(key1) == 64  # 32 bytes = 64 hex chars

    def test_sqlcipher_key_deterministic(self):
        secret = b"test-secret"
        key1 = get_sqlcipher_key("node-1", secret=secret)
        key2 = get_sqlcipher_key("node-1", secret=secret)
        assert key1 == key2

    def test_key_derivation(self):
        key = _derive_key(b"secret", b"salt", length=32)
        assert len(key) == 32
        # Same inputs = same output
        assert _derive_key(b"secret", b"salt", length=32) == key
        # Different inputs = different output
        assert _derive_key(b"secret", b"salt2", length=32) != key


# =========================================================================
# Layer 4: Sync Payload Encryption Tests
# =========================================================================

class TestSyncPayloadEncryption:
    SECRET = b"test-sync-encryption-secret-key!"

    def test_encrypt_decrypt_roundtrip(self):
        payload = {"table": "observations", "rows": [{"id": "obs-1", "title": "test"}]}
        encrypted = encrypt_sync_payload(payload, secret=self.SECRET)
        decrypted = decrypt_sync_payload(encrypted, secret=self.SECRET)
        assert decrypted == payload

    def test_encrypted_payload_has_required_fields(self):
        encrypted = encrypt_sync_payload({"test": True}, secret=self.SECRET)
        assert "ct" in encrypted
        assert "hmac" in encrypted
        assert "v" in encrypted

    def test_tampered_ciphertext_fails(self):
        encrypted = encrypt_sync_payload({"test": True}, secret=self.SECRET)
        # Tamper with ciphertext
        encrypted["ct"] = encrypted["ct"][:10] + "AAAA" + encrypted["ct"][14:]
        with pytest.raises(ValueError, match="HMAC verification failed|tampered"):
            decrypt_sync_payload(encrypted, secret=self.SECRET)

    def test_tampered_hmac_fails(self):
        encrypted = encrypt_sync_payload({"test": True}, secret=self.SECRET)
        encrypted["hmac"] = "0" * 64
        with pytest.raises(ValueError, match="HMAC verification failed|tampered"):
            decrypt_sync_payload(encrypted, secret=self.SECRET)

    def test_wrong_key_fails(self):
        encrypted = encrypt_sync_payload({"test": True}, secret=self.SECRET)
        with pytest.raises((ValueError, Exception)):
            decrypt_sync_payload(encrypted, secret=b"wrong-key-material")

    def test_large_payload(self):
        """Test encryption of a realistically-sized sync batch."""
        payload = {
            "table": "observations",
            "rows": [{"id": f"obs-{i}", "content": "x" * 500} for i in range(100)],
        }
        encrypted = encrypt_sync_payload(payload, secret=self.SECRET)
        decrypted = decrypt_sync_payload(encrypted, secret=self.SECRET)
        assert len(decrypted["rows"]) == 100

    def test_empty_payload(self):
        encrypted = encrypt_sync_payload({}, secret=self.SECRET)
        decrypted = decrypt_sync_payload(encrypted, secret=self.SECRET)
        assert decrypted == {}

    def test_unknown_version_raises(self):
        with pytest.raises(ValueError, match="Unknown encryption version"):
            decrypt_sync_payload({"v": 99, "ct": "", "hmac": ""}, secret=self.SECRET)


# =========================================================================
# Layer 4b: ML-KEM Post-Quantum Key Encapsulation Tests
# =========================================================================

@pytest.mark.skipif(not MLKEM_AVAILABLE, reason="kyber-py not installed")
class TestMLKEMEncryption:
    SECRET = b"test-mlkem-encryption-secret-key!"

    def test_keygen_returns_correct_sizes(self):
        ek, dk = generate_mlkem_keypair()
        assert len(ek) == 1184  # ML-KEM-768 encapsulation key
        assert len(dk) == 2400  # ML-KEM-768 decapsulation key

    def test_v2_encrypt_decrypt_roundtrip(self):
        ek, dk = generate_mlkem_keypair()
        payload = {"table": "observations", "rows": [{"id": "obs-1"}]}
        encrypted = encrypt_sync_payload(payload, secret=self.SECRET, peer_ek=ek)
        assert encrypted["v"] == 2
        assert "kem_ct" in encrypted
        decrypted = decrypt_sync_payload(encrypted, secret=self.SECRET, dk=dk)
        assert decrypted == payload

    def test_v2_has_required_fields(self):
        ek, _ = generate_mlkem_keypair()
        encrypted = encrypt_sync_payload({"test": True}, secret=self.SECRET, peer_ek=ek)
        assert encrypted["v"] == 2
        assert "ct" in encrypted
        assert "nonce" in encrypted
        assert "kem_ct" in encrypted
        assert "hmac" in encrypted

    def test_v2_tampered_ciphertext_fails(self):
        ek, dk = generate_mlkem_keypair()
        encrypted = encrypt_sync_payload({"test": True}, secret=self.SECRET, peer_ek=ek)
        encrypted["ct"] = encrypted["ct"][:10] + "AAAA" + encrypted["ct"][14:]
        with pytest.raises(ValueError, match="HMAC verification failed|tampered"):
            decrypt_sync_payload(encrypted, secret=self.SECRET, dk=dk)

    def test_v2_tampered_kem_ct_fails(self):
        ek, dk = generate_mlkem_keypair()
        encrypted = encrypt_sync_payload({"test": True}, secret=self.SECRET, peer_ek=ek)
        encrypted["kem_ct"] = encrypted["kem_ct"][:10] + "BBBB" + encrypted["kem_ct"][14:]
        with pytest.raises(ValueError, match="HMAC verification failed|tampered"):
            decrypt_sync_payload(encrypted, secret=self.SECRET, dk=dk)

    def test_v2_wrong_dk_fails(self):
        ek1, _ = generate_mlkem_keypair()
        _, dk2 = generate_mlkem_keypair()  # Different keypair
        encrypted = encrypt_sync_payload({"test": True}, secret=self.SECRET, peer_ek=ek1)
        # HMAC passes (it uses cluster secret), but decapsulation gives wrong key
        with pytest.raises(Exception):
            decrypt_sync_payload(encrypted, secret=self.SECRET, dk=dk2)

    def test_v2_large_payload(self):
        ek, dk = generate_mlkem_keypair()
        payload = {
            "table": "observations",
            "rows": [{"id": f"obs-{i}", "content": "x" * 500} for i in range(100)],
        }
        encrypted = encrypt_sync_payload(payload, secret=self.SECRET, peer_ek=ek)
        decrypted = decrypt_sync_payload(encrypted, secret=self.SECRET, dk=dk)
        assert len(decrypted["rows"]) == 100

    def test_v2_each_encryption_unique(self):
        """Each v2 encryption uses a fresh KEM encapsulation (ephemeral key)."""
        ek, dk = generate_mlkem_keypair()
        e1 = encrypt_sync_payload({"a": 1}, secret=self.SECRET, peer_ek=ek)
        e2 = encrypt_sync_payload({"a": 1}, secret=self.SECRET, peer_ek=ek)
        assert e1["kem_ct"] != e2["kem_ct"]  # Different ephemeral keys
        assert e1["ct"] != e2["ct"]  # Different ciphertexts

    def test_v2_without_dk_raises(self):
        ek, _ = generate_mlkem_keypair()
        encrypted = encrypt_sync_payload({"test": True}, secret=self.SECRET, peer_ek=ek)
        with pytest.raises(ValueError, match="decapsulation key"):
            decrypt_sync_payload(encrypted, secret=self.SECRET)  # No dk

    def test_save_load_keys(self, tmp_path):
        ek, dk = generate_mlkem_keypair()
        save_mlkem_keys(ek, dk, keys_dir=str(tmp_path))
        loaded_ek, loaded_dk = load_mlkem_keys(keys_dir=str(tmp_path))
        assert ek == loaded_ek
        assert dk == loaded_dk

    def test_fallback_to_v1_without_peer_ek(self):
        """Without peer_ek, encrypt_sync_payload falls back to v1."""
        encrypted = encrypt_sync_payload({"test": True}, secret=self.SECRET)
        assert encrypted["v"] == 1  # Not v2
        assert "kem_ct" not in encrypted
