"""Tests for receipt and chain verification."""

from __future__ import annotations

import base64

import cbor2
import httpx
import pytest
import respx
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.ec import SECP256R1
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey

from mailio_receipt import create_event, create_receipt
from mailio_receipt.canonicalize import canonicalize, compute_digest
from mailio_receipt.models import Actors, Intent, IntentKind, MailioAuditEvent, MailioReceipt, Signature, SignedMetadata
from mailio_receipt.verification import (
    WebAuthnPublicKey,
    _parse_cose_key,
    _verify_webauthn_signature,
    query_backend_public_keys,
)

DID_SIGNER = "did:web:mio.mailio.io:0xc3b00d0d230ac27855db6c0ff2d66d32bccf7ac7#master"
DID_DOC_URL = "https://mio.mailio.io/0xc3b00d0d230ac27855db6c0ff2d66d32bccf7ac7/did.json"
EXPECTED_PUBLIC_KEY_X = "PcrowBY5Ynz8P0tO0gq58Qa7AyzHMd9cx6TlKzQuq7s"

DID_DOCUMENT = {
    "@context": [
        "https://www.w3.org/ns/did/v1",
        "https://w3id.org/security/suites/ed25519-2020/v1",
        "https://w3id.org/security/suites/x25519-2019/v1",
    ],
    "id": "did:mailio:0xc3b00d0d230ac27855db6c0ff2d66d32bccf7ac7",
    "authentication": [
        "did:mailio:0xc3b00d0d230ac27855db6c0ff2d66d32bccf7ac7#master"
    ],
    "verificationMethod": [
        {
            "id": "did:mailio:0xc3b00d0d230ac27855db6c0ff2d66d32bccf7ac7#master",
            "type": "JsonWebKey2020",
            "controller": "did:mailio:0xc3b00d0d230ac27855db6c0ff2d66d32bccf7ac7",
            "publicKeyJwk": {
                "crv": "Ed25519",
                "kty": "OKP",
                "x": EXPECTED_PUBLIC_KEY_X,
            },
        }
    ],
}


class TestQueryBackendPublicKeys:
    @pytest.mark.asyncio
    @respx.mock
    async def test_query_did_key_returns_correct_ed25519_public_key(self) -> None:
        """Build a receipt signed by a DID signer and verify query_backend_public_keys
        resolves the correct Ed25519 public key."""
        respx.get(DID_DOC_URL).mock(return_value=httpx.Response(200, json=DID_DOCUMENT))

        private_key = Ed25519PrivateKey.generate()
        event = create_event(
            flow_id="c03f0ebc-a387-42c8-8057-3af58801b794",
            actors=Actors(issuer=DID_SIGNER, subject=DID_SIGNER),
            intent=Intent(
                kind=IntentKind.ACTION_PLANNED,
                action="test.action",
                category="test",
            ),
        )
        receipt = create_receipt(
            event=event,
            private_key=private_key,
            signer_did=DID_SIGNER,
            role="agent",
        )

        result = await query_backend_public_keys(receipt)

        assert DID_SIGNER in result
        public_key = result[DID_SIGNER]
        assert isinstance(public_key, Ed25519PublicKey)

        raw_bytes = public_key.public_bytes_raw()
        x_b64 = base64.urlsafe_b64encode(raw_bytes).rstrip(b"=").decode()
        assert x_b64 == EXPECTED_PUBLIC_KEY_X


def _ec_key_to_cose_b64(public_key: ec.EllipticCurvePublicKey) -> str:
    """Encode an EC P-256 public key as base64-encoded COSE key map."""
    numbers = public_key.public_numbers()
    cose_map = {
        1: 2,       # kty: EC2
        3: -7,      # alg: ES256
        -1: 1,      # crv: P-256
        -2: numbers.x.to_bytes(32, "big"),
        -3: numbers.y.to_bytes(32, "big"),
    }
    return base64.b64encode(cbor2.dumps(cose_map)).decode()


class TestWebAuthnCoseVerification:
    """Tests for COSE key parsing and WebAuthn signature verification."""

    def test_parse_cose_key_roundtrip(self) -> None:
        """Encode an EC key as COSE, parse it back, and verify coordinates match."""
        private_key = ec.generate_private_key(SECP256R1())
        original_pub = private_key.public_key()
        cose_b64 = _ec_key_to_cose_b64(original_pub)

        parsed_pub = _parse_cose_key(cose_b64)

        assert original_pub.public_numbers().x == parsed_pub.public_numbers().x
        assert original_pub.public_numbers().y == parsed_pub.public_numbers().y

    def test_parse_cose_key_rejects_wrong_kty(self) -> None:
        """COSE key with non-EC2 kty should be rejected."""
        cose_map = {1: 1, 3: -7, -1: 1, -2: b"\x00" * 32, -3: b"\x00" * 32}
        b64 = base64.b64encode(cbor2.dumps(cose_map)).decode()
        with pytest.raises(ValueError, match="Unsupported COSE key type"):
            _parse_cose_key(b64)

    def test_parse_cose_key_rejects_wrong_curve(self) -> None:
        """COSE key with non-P-256 curve should be rejected."""
        cose_map = {1: 2, 3: -7, -1: 2, -2: b"\x00" * 32, -3: b"\x00" * 32}
        b64 = base64.b64encode(cbor2.dumps(cose_map)).decode()
        with pytest.raises(ValueError, match="Unsupported COSE curve"):
            _parse_cose_key(b64)

    def test_verify_webauthn_assertion_synthetic(self) -> None:
        """Sign authenticatorData || SHA256(clientDataJSON), verify via _verify_webauthn_signature."""
        import hashlib
        import json

        private_key = ec.generate_private_key(SECP256R1())
        public_key = private_key.public_key()

        # Build synthetic WebAuthn assertion data
        authenticator_data = b"\x49\x96\x0d\xe5" + b"\x00" * 33  # 37 bytes rpIdHash + flags + counter
        challenge = "dGVzdC1jaGFsbGVuZ2U"  # base64url of "test-challenge"
        client_data = json.dumps({
            "type": "webauthn.get",
            "challenge": challenge,
            "origin": "http://localhost",
            "crossOrigin": False,
        }).encode("utf-8")

        # Authenticator signs: authenticatorData || SHA256(clientDataJSON)
        client_data_hash = hashlib.sha256(client_data).digest()
        signed_data = authenticator_data + client_data_hash
        signature_der = private_key.sign(signed_data, ec.ECDSA(hashes.SHA256()))

        sig_b64 = base64.urlsafe_b64encode(signature_der).rstrip(b"=").decode()
        auth_data_b64 = base64.urlsafe_b64encode(authenticator_data).rstrip(b"=").decode()
        client_data_b64 = base64.urlsafe_b64encode(client_data).rstrip(b"=").decode()

        assert _verify_webauthn_signature(sig_b64, public_key, auth_data_b64, client_data_b64) is True

    def test_verify_webauthn_assertion_wrong_client_data(self) -> None:
        """Verification must fail when clientDataJSON doesn't match."""
        import hashlib
        import json

        private_key = ec.generate_private_key(SECP256R1())
        public_key = private_key.public_key()

        authenticator_data = b"\x49\x96\x0d\xe5" + b"\x00" * 33
        client_data = json.dumps({
            "type": "webauthn.get",
            "challenge": "original",
            "origin": "http://localhost",
            "crossOrigin": False,
        }).encode("utf-8")

        client_data_hash = hashlib.sha256(client_data).digest()
        signed_data = authenticator_data + client_data_hash
        signature_der = private_key.sign(signed_data, ec.ECDSA(hashes.SHA256()))

        sig_b64 = base64.urlsafe_b64encode(signature_der).rstrip(b"=").decode()
        auth_data_b64 = base64.urlsafe_b64encode(authenticator_data).rstrip(b"=").decode()

        # Tampered clientDataJSON with different challenge
        tampered_client_data = json.dumps({
            "type": "webauthn.get",
            "challenge": "tampered",
            "origin": "http://localhost",
            "crossOrigin": False,
        }).encode("utf-8")
        tampered_b64 = base64.urlsafe_b64encode(tampered_client_data).rstrip(b"=").decode()

        assert _verify_webauthn_signature(sig_b64, public_key, auth_data_b64, tampered_b64) is False

    def test_verify_webauthn_assertion_wrong_key(self) -> None:
        """Verification must fail with a different public key."""
        import hashlib
        import json

        private_key = ec.generate_private_key(SECP256R1())
        wrong_key = ec.generate_private_key(SECP256R1()).public_key()

        authenticator_data = b"\x49\x96\x0d\xe5" + b"\x00" * 33
        client_data = json.dumps({
            "type": "webauthn.get",
            "challenge": "test",
            "origin": "http://localhost",
            "crossOrigin": False,
        }).encode("utf-8")

        client_data_hash = hashlib.sha256(client_data).digest()
        signed_data = authenticator_data + client_data_hash
        signature_der = private_key.sign(signed_data, ec.ECDSA(hashes.SHA256()))

        sig_b64 = base64.urlsafe_b64encode(signature_der).rstrip(b"=").decode()
        auth_data_b64 = base64.urlsafe_b64encode(authenticator_data).rstrip(b"=").decode()
        client_data_b64 = base64.urlsafe_b64encode(client_data).rstrip(b"=").decode()

        assert _verify_webauthn_signature(sig_b64, wrong_key, auth_data_b64, client_data_b64) is False

    def test_full_cose_key_and_webauthn_roundtrip(self) -> None:
        """End-to-end: generate key, encode as COSE, parse back, sign WebAuthn assertion, verify."""
        import hashlib
        import json

        private_key = ec.generate_private_key(SECP256R1())
        public_key = private_key.public_key()

        # Encode public key as COSE (simulating what the backend stores)
        cose_b64 = _ec_key_to_cose_b64(public_key)
        parsed_key = _parse_cose_key(cose_b64)

        # Compute event digest (simulating what the server sends as challenge)
        event = create_event(
            flow_id="c03f0ebc-a387-42c8-8057-3af58801b794",
            actors=Actors(issuer="user@mail.io", subject="agent-credential-id"),
            intent=Intent(
                kind=IntentKind.APPROVAL_GRANTED,
                action="test.action",
                category="test",
            ),
        )
        digest = compute_digest(event.to_dict())

        # Build WebAuthn assertion (simulating browser authenticator)
        authenticator_data = b"\x49\x96\x0d\xe5" + b"\x00" * 33
        client_data = json.dumps({
            "type": "webauthn.get",
            "challenge": digest,
            "origin": "http://localhost",
            "crossOrigin": False,
        }).encode("utf-8")

        client_data_hash = hashlib.sha256(client_data).digest()
        signed_data = authenticator_data + client_data_hash
        signature_der = private_key.sign(signed_data, ec.ECDSA(hashes.SHA256()))

        sig_b64 = base64.urlsafe_b64encode(signature_der).rstrip(b"=").decode()
        auth_data_b64 = base64.urlsafe_b64encode(authenticator_data).rstrip(b"=").decode()
        client_data_b64 = base64.urlsafe_b64encode(client_data).rstrip(b"=").decode()

        # Verify using the parsed COSE key
        assert _verify_webauthn_signature(sig_b64, parsed_key, auth_data_b64, client_data_b64) is True

    @pytest.mark.asyncio
    @respx.mock
    async def test_query_webauthn_key_returns_wrapped_type(self) -> None:
        """query_backend_public_keys wraps webauthn keys in WebAuthnPublicKey."""
        private_key = ec.generate_private_key(SECP256R1())
        public_key = private_key.public_key()
        cose_b64 = _ec_key_to_cose_b64(public_key)

        signer = "webauthn-credential-abc123"
        webauthn_url = f"https://mio.mailiomail.com/api/v1/receipts/webauthn/public_key/{signer}"
        respx.get(webauthn_url).mock(
            return_value=httpx.Response(200, json={"publicKey": cose_b64})
        )

        # Build a minimal receipt with the webauthn signer
        canonical_event = canonicalize(
            create_event(
                flow_id="c03f0ebc-a387-42c8-8057-3af58801b794",
                actors=Actors(issuer="user@mail.io", subject=signer),
                intent=Intent(
                    kind=IntentKind.APPROVAL_GRANTED,
                    action="test.action",
                    category="test",
                ),
            ).to_dict()
        )
        signature_der = private_key.sign(canonical_event, ec.ECDSA(hashes.SHA256()))
        sig_b64 = base64.urlsafe_b64encode(signature_der).rstrip(b"=").decode()

        event = create_event(
            flow_id="c03f0ebc-a387-42c8-8057-3af58801b794",
            actors=Actors(issuer="user@mail.io", subject=signer),
            intent=Intent(
                kind=IntentKind.APPROVAL_GRANTED,
                action="test.action",
                category="test",
            ),
        )
        receipt = MailioReceipt(
            event=event,
            event_digest=compute_digest(event.to_dict()),
            signatures=(
                Signature(
                    role="user",
                    signer=signer,
                    alg="ES256",
                    signed=SignedMetadata(
                        canonicalization="RFC8785",
                        digest=compute_digest(event.to_dict()),
                        target="/event",
                    ),
                    jws=sig_b64,
                ),
            ),
        )

        result = await query_backend_public_keys(receipt)

        assert signer in result
        key = result[signer]
        assert isinstance(key, WebAuthnPublicKey)
        assert key.key.public_numbers().x == public_key.public_numbers().x


class TestKnownWebAuthnAssertion:
    """Verify a real WebAuthn assertion from a known receipt against a known COSE public key."""

    COSE_PUBLIC_KEY_B64 = (
        "pQECAyYgASFYIKtxrXJVma2LVBmNi5dXL/EIUZxXdWDV16PEbGUac06y"
        "Ilggds8F5U8/ZdsVvEenTDspUzMoG2RQSX8sxxzHpWqN/r8="
    )

    EVENT_DICT: dict = {
        "type": "MailioAuditEvent",
        "v": "v1",
        "id": "0e8faf1c-3527-4190-9556-25d0c7ccccc8",
        "flow_id": "88e6cfcb-e2d3-4319-bc35-3e03f2ddf2b2",
        "ts": "2026-02-13T17:18:46.361Z",
        "actors": {
            "issuer": "urn:mailio:igordemo-local-69676f72406d61696c2e696f",
            "subject": "did:web:mio.mailiomail.com:0xb3e49d751ba7d6647a720a52060ceebe42524704",
        },
        "intent": {
            "kind": "trigger_watch_condition",
            "action": "price_above",
            "category": "stock_price",
            "scope": {
                "ticker": "MIO",
                "direction": "price_above",
                "threshold": 77,
            },
        },
    }

    EXPECTED_DIGEST = "USSjifDkdf7dvNNkUcLlWHSC1aec8ebruvDINnmcGl4"

    SIGNATURE_B64 = (
        "MEUCIGl2Lc8n2gVlctBquDppN58YPdqNNw7r4sQuRyWBki55"
        "AiEAzMjqiHuhvqey501twLU16S09ZWI9UzPfZ9OA6oBefcE"
    )
    AUTHENTICATOR_DATA_B64 = "SZYN5YgOjGh0NBcPZHZgW4_krrmihjLHmVzzuoMdl2MdAAAAAA"
    CLIENT_DATA_JSON_B64 = (
        "eyJ0eXBlIjoid2ViYXV0aG4uZ2V0IiwiY2hhbGxlbmdlIjoiVVNTamlm"
        "RGtkZjdkdk5Oa1VjTGxXSFNDMWFlYzhlYnJ1dkRJTm5tY0dsNCIsIm9y"
        "aWdpbiI6Imh0dHA6Ly9sb2NhbGhvc3Q6ODg4OCIsImNyb3NzT3JpZ2lu"
        "IjpmYWxzZX0"
    )

    SIGNER = "urn:mailio:igordemo-local-69676f72406d61696c2e696f#webauthn-cred:EBACb3n-2F1HOkqIuCRRyA"

    def test_parse_known_cose_public_key(self) -> None:
        """The COSE public key should parse into a valid EC P-256 key."""
        public_key = _parse_cose_key(self.COSE_PUBLIC_KEY_B64)
        numbers = public_key.public_numbers()
        assert isinstance(numbers.curve, SECP256R1)
        assert numbers.x.bit_length() <= 256
        assert numbers.y.bit_length() <= 256

    def test_event_digest_matches(self) -> None:
        """The CBOR-canonicalized event should produce the expected base64url digest."""
        event = MailioAuditEvent.from_dict(self.EVENT_DICT)
        digest = compute_digest(event.to_dict())
        assert digest == self.EXPECTED_DIGEST

    def test_challenge_matches_digest(self) -> None:
        """The WebAuthn challenge in clientDataJSON must equal the event digest."""
        import json

        from mailio_receipt.verification import _b64url_decode

        client_data = json.loads(_b64url_decode(self.CLIENT_DATA_JSON_B64))
        assert client_data["type"] == "webauthn.get"
        assert client_data["challenge"] == self.EXPECTED_DIGEST

    def test_verify_webauthn_assertion(self) -> None:
        """The full WebAuthn assertion must verify against the known public key."""
        public_key = _parse_cose_key(self.COSE_PUBLIC_KEY_B64)

        assert _verify_webauthn_signature(
            self.SIGNATURE_B64,
            public_key,
            self.AUTHENTICATOR_DATA_B64,
            self.CLIENT_DATA_JSON_B64,
        ) is True

    def test_verify_rejects_wrong_key(self) -> None:
        """Verification must fail with a different public key."""
        wrong_key = ec.generate_private_key(SECP256R1()).public_key()

        assert _verify_webauthn_signature(
            self.SIGNATURE_B64,
            wrong_key,
            self.AUTHENTICATOR_DATA_B64,
            self.CLIENT_DATA_JSON_B64,
        ) is False

    def test_verify_rejects_tampered_authenticator_data(self) -> None:
        """Verification must fail when authenticator data is altered."""
        public_key = _parse_cose_key(self.COSE_PUBLIC_KEY_B64)

        assert _verify_webauthn_signature(
            self.SIGNATURE_B64,
            public_key,
            "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
            self.CLIENT_DATA_JSON_B64,
        ) is False

    def test_full_receipt_roundtrip(self) -> None:
        """Parse the full receipt from JSON and verify the event digest and model integrity."""
        receipt = MailioReceipt.from_dict({
            "type": "MailioReceipt",
            "v": "v1",
            "event": self.EVENT_DICT,
            "event_digest": self.EXPECTED_DIGEST,
            "signatures": [
                {
                    "role": "user",
                    "signer": self.SIGNER,
                    "alg": "ES256",
                    "signed": {
                        "canonicalization": "RFC8785",
                        "digest": self.EXPECTED_DIGEST,
                        "target": "/event",
                    },
                    "jws": self.SIGNATURE_B64,
                    "authenticator_data": self.AUTHENTICATOR_DATA_B64,
                    "client_data_json": self.CLIENT_DATA_JSON_B64,
                },
            ],
        })

        # Verify digest matches
        assert compute_digest(receipt.event.to_dict()) == receipt.event_digest

        # Verify signature fields are preserved
        sig = receipt.signatures[0]
        assert sig.authenticator_data == self.AUTHENTICATOR_DATA_B64
        assert sig.client_data_json == self.CLIENT_DATA_JSON_B64

        # Verify the signature
        public_key = _parse_cose_key(self.COSE_PUBLIC_KEY_B64)
        assert _verify_webauthn_signature(
            sig.jws, public_key, sig.authenticator_data, sig.client_data_json
        ) is True
