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

from __future__ import annotations

from typing import TYPE_CHECKING, Any

import pytest

if TYPE_CHECKING:
    from cryptography.hazmat.primitives.asymmetric.ec import (
        EllipticCurvePrivateKey,
        EllipticCurvePublicKey,
    )

from mailio_receipt import (
    MailioReceipt,
    Actors,
    Intent,
    IntentKind,
    chain_from,
    create_event,
    create_receipt,
)

import base64
from mailio_receipt.canonicalize import compute_digest

from mailio_receipt.errors import (
    InvalidSignatureError,
    MissingPredecessorError,
)
from mailio_receipt.verification import verify_chain, verify_receipt


class MockDIDResolver:
    """Mock DID resolver for testing."""

    def __init__(self, keys: dict[str, EllipticCurvePublicKey]) -> None:
        self.keys = keys

    async def resolve(self, did: str) -> dict[str, Any]:
        return {"id": did}

    async def get_verification_key(
        self, did: str, key_id: str | None = None
    ) -> EllipticCurvePublicKey:
        if did not in self.keys:
            raise KeyError(f"DID not found: {did}")
        return self.keys[did]

    async def close(self) -> None:
        pass


class TestVerifyReceipt:
    """Tests for single receipt verification."""

    @pytest.mark.asyncio
    async def test_verify_valid_receipt(
        self,
        agent_private_key: EllipticCurvePrivateKey,
        agent_public_key: EllipticCurvePublicKey,
        sample_actors: Actors,
        sample_intent: Intent,
        sample_flow_id: str,
        agent_did: str,
    ) -> None:
        """Test that valid receipt verifies successfully."""
        # Create receipt
        event = create_event(
            flow_id=sample_flow_id,
            actors=sample_actors,
            intent=sample_intent,
        )
        receipt = create_receipt(
            event=event,
            private_key=agent_private_key,
            signer_did=agent_did,
        )

        # Verify with mock resolver
        resolver = MockDIDResolver({agent_did: agent_public_key})
        result = await verify_receipt(receipt, resolver=resolver)

        assert result.valid is True
        assert len(result.signers) == 1
        assert result.signers[0].did == agent_did
        assert result.signers[0].role == "agent"
        assert result.error is None

    @pytest.mark.asyncio
    async def test_verify_event_digest_mismatch(
        self,
        agent_private_key: EllipticCurvePrivateKey,
        agent_public_key: EllipticCurvePublicKey,
        sample_actors: Actors,
        sample_intent: Intent,
        sample_flow_id: str,
        agent_did: str,
    ) -> None:
        """Test that tampered event digest is detected."""
        # Create receipt
        event = create_event(
            flow_id=sample_flow_id,
            actors=sample_actors,
            intent=sample_intent,
        )
        receipt = create_receipt(
            event=event,
            private_key=agent_private_key,
            signer_did=agent_did,
        )

        # Tamper with event digest using object.__setattr__ since dataclass is frozen
        tampered_receipt = receipt.__class__(
            type=receipt.type,
            v=receipt.v,
            event=receipt.event,
            event_digest="a" * 64,  # Wrong digest
            signatures=receipt.signatures,
        )

        # Verify
        resolver = MockDIDResolver({agent_did: agent_public_key})
        result = await verify_receipt(tampered_receipt, resolver=resolver)

        assert result.valid is False
        assert "digest mismatch" in result.error.lower()

    @pytest.mark.asyncio
    async def test_verify_invalid_signature(
        self,
        agent_private_key: EllipticCurvePrivateKey,
        user_public_key: EllipticCurvePublicKey,  # Wrong key
        sample_actors: Actors,
        sample_intent: Intent,
        sample_flow_id: str,
        agent_did: str,
    ) -> None:
        """Test that wrong public key fails verification."""
        # Create receipt
        event = create_event(
            flow_id=sample_flow_id,
            actors=sample_actors,
            intent=sample_intent,
        )
        receipt = create_receipt(
            event=event,
            private_key=agent_private_key,
            signer_did=agent_did,
        )

        # Verify with wrong key
        resolver = MockDIDResolver({agent_did: user_public_key})  # Wrong key!

        with pytest.raises(InvalidSignatureError):
            await verify_receipt(receipt, resolver=resolver)

    @pytest.mark.asyncio
    async def test_verify_tampered_event_content(
        self,
        agent_private_key: EllipticCurvePrivateKey,
        agent_public_key: EllipticCurvePublicKey,
        sample_actors: Actors,
        sample_flow_id: str,
        agent_did: str,
    ) -> None:
        """Test that tampered event content is detected via signature."""
        # Create receipt with one intent
        original_intent = Intent(
            kind=IntentKind.ACTION_PLANNED,
            action="invoice.send",
            category="financial",
            scope={"amount": 1000},
        )
        event = create_event(
            flow_id=sample_flow_id,
            actors=sample_actors,
            intent=original_intent,
        )
        receipt = create_receipt(
            event=event,
            private_key=agent_private_key,
            signer_did=agent_did,
        )

        # Create tampered event with different amount
        tampered_intent = Intent(
            kind=IntentKind.ACTION_PLANNED,
            action="invoice.send",
            category="financial",
            scope={"amount": 9999},  # Tampered!
        )
        tampered_event = create_event(
            flow_id=sample_flow_id,
            actors=sample_actors,
            intent=tampered_intent,
            event_id=receipt.event.id,
            timestamp=None,  # Will generate new timestamp but that's fine
        )

        # Create tampered receipt (digest will be wrong)
        tampered_receipt = receipt.__class__(
            type=receipt.type,
            v=receipt.v,
            event=tampered_event,
            event_digest=receipt.event_digest,  # Old digest
            signatures=receipt.signatures,  # Old signatures
        )

        # Verify - should fail due to digest mismatch
        resolver = MockDIDResolver({agent_did: agent_public_key})
        result = await verify_receipt(tampered_receipt, resolver=resolver)

        assert result.valid is False


class TestVerifyChain:
    """Tests for chain verification."""

    @pytest.mark.asyncio
    async def test_verify_valid_chain(
        self,
        agent_private_key: EllipticCurvePrivateKey,
        user_private_key: EllipticCurvePrivateKey,
        agent_public_key: EllipticCurvePublicKey,
        user_public_key: EllipticCurvePublicKey,
        sample_actors: Actors,
        sample_intent: Intent,
        sample_flow_id: str,
        agent_did: str,
        user_did: str,
    ) -> None:
        """Test that valid chain verifies successfully."""
        # Create chain: planned -> approved
        planned_event = create_event(
            flow_id=sample_flow_id,
            actors=sample_actors,
            intent=sample_intent,
        )
        planned_receipt = create_receipt(
            event=planned_event,
            private_key=agent_private_key,
            signer_did=agent_did,
        )

        user_actors = Actors(issuer=user_did, subject=user_did)
        approval_intent = Intent(
            kind="approval",
            action=planned_receipt.event.intent.action,
            category=planned_receipt.event.intent.category,
            scope={"approvedReceiptDigest": planned_receipt.event_digest},
        )
        approval_event = create_event(
            flow_id=sample_flow_id,
            actors=user_actors,
            intent=approval_intent,
            chain=chain_from(planned_receipt),
        )
        approval_receipt = create_receipt(
            event=approval_event,
            private_key=user_private_key,
            signer_did=user_did,
            role="user",
        )

        # Verify chain
        resolver = MockDIDResolver({agent_did: agent_public_key, user_did: user_public_key})
        result = await verify_chain([planned_receipt, approval_receipt], resolver=resolver)

        assert result.valid is True
        assert result.chain_intact is True
        assert len(result.receipt_results) == 2
        assert all(r.valid for r in result.receipt_results)

    @pytest.mark.asyncio
    async def test_verify_chain_missing_predecessor(
        self,
        agent_private_key: EllipticCurvePrivateKey,
        user_private_key: EllipticCurvePrivateKey,
        agent_public_key: EllipticCurvePublicKey,
        user_public_key: EllipticCurvePublicKey,
        sample_actors: Actors,
        sample_intent: Intent,
        sample_flow_id: str,
        agent_did: str,
        user_did: str,
    ) -> None:
        """Test that missing predecessor is detected."""
        # Create chain but only provide the approval (missing planned)
        planned_event = create_event(
            flow_id=sample_flow_id,
            actors=sample_actors,
            intent=sample_intent,
        )
        planned_receipt = create_receipt(
            event=planned_event,
            private_key=agent_private_key,
            signer_did=agent_did,
        )

        user_actors = Actors(issuer=user_did, subject=user_did)
        approval_intent = Intent(
            kind="approval",
            action=planned_receipt.event.intent.action,
            category=planned_receipt.event.intent.category,
            scope={"approvedReceiptDigest": planned_receipt.event_digest},
        )
        approval_event = create_event(
            flow_id=sample_flow_id,
            actors=user_actors,
            intent=approval_intent,
            chain=chain_from(planned_receipt),
        )
        approval_receipt = create_receipt(
            event=approval_event,
            private_key=user_private_key,
            signer_did=user_did,
            role="user",
        )

        # Verify chain with missing predecessor
        resolver = MockDIDResolver({agent_did: agent_public_key, user_did: user_public_key})

        with pytest.raises(MissingPredecessorError):
            await verify_chain([approval_receipt], resolver=resolver)  # Missing planned!

    @pytest.mark.asyncio
    async def test_verify_chain_with_invalid_receipt(
        self,
        agent_private_key: EllipticCurvePrivateKey,
        user_private_key: EllipticCurvePrivateKey,
        agent_public_key: EllipticCurvePublicKey,
        user_public_key: EllipticCurvePublicKey,
        sample_actors: Actors,
        sample_intent: Intent,
        sample_flow_id: str,
        agent_did: str,
        user_did: str,
    ) -> None:
        """Test that chain fails if any receipt is invalid."""
        # Create valid planned receipt
        planned_event = create_event(
            flow_id=sample_flow_id,
            actors=sample_actors,
            intent=sample_intent,
        )
        planned_receipt = create_receipt(
            event=planned_event,
            private_key=agent_private_key,
            signer_did=agent_did,
        )

        # Create approval but tamper with its digest
        user_actors = Actors(issuer=user_did, subject=user_did)
        approval_intent = Intent(
            kind="approval",
            action=planned_receipt.event.intent.action,
            category=planned_receipt.event.intent.category,
            scope={"approvedReceiptDigest": planned_receipt.event_digest},
        )
        approval_event = create_event(
            flow_id=sample_flow_id,
            actors=user_actors,
            intent=approval_intent,
            chain=chain_from(planned_receipt),
        )
        approval_receipt = create_receipt(
            event=approval_event,
            private_key=user_private_key,
            signer_did=user_did,
            role="user",
        )

        # Tamper with approval's event digest
        tampered_approval = approval_receipt.__class__(
            type=approval_receipt.type,
            v=approval_receipt.v,
            event=approval_receipt.event,
            event_digest="b" * 64,  # Wrong digest
            signatures=approval_receipt.signatures,
        )

        # Verify chain
        resolver = MockDIDResolver({agent_did: agent_public_key, user_did: user_public_key})
        result = await verify_chain([planned_receipt, tampered_approval], resolver=resolver)

        assert result.valid is False
        assert result.chain_intact is False

    @pytest.mark.asyncio
    async def test_verify_single_receipt_chain(
        self,
        agent_private_key: EllipticCurvePrivateKey,
        agent_public_key: EllipticCurvePublicKey,
        sample_actors: Actors,
        sample_intent: Intent,
        sample_flow_id: str,
        agent_did: str,
    ) -> None:
        """Test verification of a single receipt (no predecessors)."""
        # Create single receipt with no chain
        event = create_event(
            flow_id=sample_flow_id,
            actors=sample_actors,
            intent=sample_intent,
        )
        receipt = create_receipt(
            event=event,
            private_key=agent_private_key,
            signer_did=agent_did,
        )

        # Verify
        resolver = MockDIDResolver({agent_did: agent_public_key})
        result = await verify_chain([receipt], resolver=resolver)

        assert result.valid is True
        assert result.chain_intact is True

    @pytest.mark.asyncio
    async def test_verify_webauthn_signature(self):
        data = {
            "type": "MailioReceipt",
            "v": "v1",
            "event": {
                "type": "MailioAuditEvent",
                "v": "v1",
                "id": "fd769520-42d1-4b84-bc4a-42d32df039bc",
                "flowId": "2eeb0bdc-b885-4547-8514-d4e3f33206cb",
                "ts": "2026-02-11T05:38:19.862Z",
                "actors": {
                    "issuer": "test2@mail.io",
                    "subject": "igordemo-local-69676f72406d61696c2e696f"
                },
                "intent": {
                    "kind": "trigger_watch_condition",
                    "action": "price_above",
                    "category": "stock_price",
                    "scope": {
                        "ticker": "MIO",
                        "direction": "price_above",
                        "threshold": 23
                    },
                },
            },
            "event_digest": "AyMOTKC0GZsk74MsxUJcfcqA8LUfZ1NZ9nCwBpvc5YM",
            "signatures": [
                {
                    "role": "user",
                    "signer": "igordemo-local-69676f72406d61696c2e696f",
                    "alg": "ES256",
                    "signed": {
                        "canonicalization": "RFC8785",
                        "digest": "AyMOTKC0GZsk74MsxUJcfcqA8LUfZ1NZ9nCwBpvc5YM",
                        "target": "/event",
                    },
                    "jws": "MEUCIBO_ypa3TetBHcfVnO1-Yh4xbDzNYmiL_1AP9pO8hFOIAiEApfefA8OL2ydn8vVrZm-weSi2fPI2DVKwgRu_APHV81c",
                },
            ],
        }

        receipt = MailioReceipt.from_dict(data)
        try:
            result = await verify_receipt(receipt)
            print(result)
        except Exception as e:
            print(f"Verification failed with error: {e}")
            raise

        



