"""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
from mailio_receipt.models import Actors, Intent, IntentKind, 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_der_signature(self) -> None:
        """Sign payload with EC key (producing DER), verify via _verify_webauthn_signature."""
        private_key = ec.generate_private_key(SECP256R1())
        public_key = private_key.public_key()
        payload = b'{"hello":"world"}'

        # Sign directly — produces DER-encoded signature (like a WebAuthn client)
        signature_der = private_key.sign(payload, ec.ECDSA(hashes.SHA256()))
        sig_b64 = base64.urlsafe_b64encode(signature_der).rstrip(b"=").decode()

        assert _verify_webauthn_signature(sig_b64, payload, public_key) is True

    def test_verify_webauthn_signature_wrong_payload(self) -> None:
        """Verification must fail when payload doesn't match."""
        private_key = ec.generate_private_key(SECP256R1())
        public_key = private_key.public_key()

        signature_der = private_key.sign(b"original", ec.ECDSA(hashes.SHA256()))
        sig_b64 = base64.urlsafe_b64encode(signature_der).rstrip(b"=").decode()

        assert _verify_webauthn_signature(sig_b64, b"tampered", public_key) is False

    def test_verify_webauthn_signature_wrong_key(self) -> None:
        """Verification must fail with a different public key."""
        private_key = ec.generate_private_key(SECP256R1())
        wrong_key = ec.generate_private_key(SECP256R1()).public_key()
        payload = b'{"test":"data"}'

        signature_der = private_key.sign(payload, ec.ECDSA(hashes.SHA256()))
        sig_b64 = base64.urlsafe_b64encode(signature_der).rstrip(b"=").decode()

        assert _verify_webauthn_signature(sig_b64, payload, wrong_key) is False

    def test_full_cose_key_and_signature_roundtrip(self) -> None:
        """End-to-end: generate key, encode as COSE, parse back, sign, verify."""
        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)

        # Parse it back (simulating what _fetch_webauthn_key does)
        parsed_key = _parse_cose_key(cose_b64)

        # Sign a canonical event payload (simulating what the WebAuthn client does)
        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",
            ),
        )
        canonical_event = canonicalize(event.to_dict())
        signature_der = private_key.sign(canonical_event, ec.ECDSA(hashes.SHA256()))
        sig_b64 = base64.urlsafe_b64encode(signature_der).rstrip(b"=").decode()

        # Verify using the parsed COSE key
        assert _verify_webauthn_signature(sig_b64, canonical_event, parsed_key) 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()

        from mailio_receipt.canonicalize import compute_digest
        from mailio_receipt.models import MailioReceipt

        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
