from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from mailio_receipt import chain_from, create_receipt, verify_receipt
from mailio_receipt.models import Actors, Intent, create_event
import pytest
import base64
import respx

class TestIgor:
    """Tests for Igor."""

    def _ed25519_public_key_to_jwk(self, public_key: Ed25519PublicKey) -> dict:
        """Convert Ed25519 public key to JWK."""
        x_bytes = public_key.public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw,
        )
        return {
            "kty": "OKP",
            "crv": "Ed25519",
            "x": base64.urlsafe_b64encode(x_bytes).decode("ascii").rstrip("="),
        }

    @pytest.mark.asyncio
    @respx.mock
    async def test_create_receipt(self, 
        ed25519_agent_private_key: Ed25519PrivateKey,
        ed25519_agent_public_key: Ed25519PublicKey,
        sample_actors: Actors,
        sample_intent: Intent,
        sample_flow_id: str,
        agent_did: str,
    ) -> None:
        """Test that create_receipt creates a valid receipt."""
        event = create_event(
            flow_id=sample_flow_id,
            actors=sample_actors,
            intent=sample_intent,
        )
        receipt = create_receipt(
            event=event,
            private_key=ed25519_agent_private_key,
            signer_did=agent_did,
        )
        print(receipt.to_json())

        did_doc = {
            "id": agent_did,
            "@context": ["https://www.w3.org/ns/did/v1"],
            "verificationMethod": [
                {
                    "id": agent_did,
                    "type": "JsonWebKey2020",
                    "controller": agent_did,
                    "publicKeyJwk": self._ed25519_public_key_to_jwk(
                        ed25519_agent_public_key
                    ),
                }
            ],
        }
        respx.get("https://agent.example.com/.well-known/did.json").respond(json=did_doc)

        result = await verify_receipt(receipt)
        for signer in result.signers:
            print(f"Verified signer: {signer.did} with role {signer.role} and key ID {signer.key_id}")
        assert result.valid, f"Verification failed: {result.error}"

   
    @pytest.mark.asyncio
    @respx.mock
    async def test_verify_chained_receipts(
        self,
        ed25519_agent_private_key: Ed25519PrivateKey,
        ed25519_agent_public_key: Ed25519PublicKey,
        ed25519_user_private_key: Ed25519PrivateKey,
        ed25519_user_public_key: Ed25519PublicKey,
        sample_actors: Actors,
        sample_intent: Intent,
        sample_flow_id: str,
        agent_did: str,
        user_did: str,
    ) -> None:
        event = create_event(
            flow_id=sample_flow_id,
            actors=sample_actors,
            intent=sample_intent,
        )
        planned_receipt = create_receipt(
            event=event,
            private_key=ed25519_agent_private_key,
            signer_did=agent_did,
        )
        # Step 2: Create a chained follow-up receipt
        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=ed25519_user_private_key,
            signer_did=user_did,
            role="user",
        )

        # Step 3: Create another chained receipt
        revocation_intent = Intent(
            kind="revocation",
            action=planned_receipt.event.intent.action,
            category=planned_receipt.event.intent.category,
            scope={
                "revokedReceiptDigest": approval_receipt.event_digest,
            },
        )
        revocation_event = create_event(
            flow_id=sample_flow_id,
            actors=user_actors,
            intent=revocation_intent,
            chain=chain_from(approval_receipt),
        )
        revocation_receipt = create_receipt(
            event=revocation_event,
            private_key=ed25519_user_private_key,
            signer_did=user_did,
            role="user",
        )

        # valida signatures from the complete chain, starting with initial receipt, then approval and finally revocation
        did_doc_agent = {
            "id": agent_did,
            "@context": ["https://www.w3.org/ns/did/v1"],
            "verificationMethod": [
                {
                    "id": agent_did,
                    "type": "JsonWebKey2020",
                    "controller": agent_did,
                    "publicKeyJwk": self._ed25519_public_key_to_jwk(
                        ed25519_agent_public_key
                    ),
                }
            ],
        }
        did_doc_user = {
            "id": user_did,
            "@context": ["https://www.w3.org/ns/did/v1"],
            "verificationMethod": [ 
                {
                    "id": user_did,
                    "type": "JsonWebKey2020",
                    "controller": user_did,
                    "publicKeyJwk": self._ed25519_public_key_to_jwk(
                        ed25519_user_public_key
                    ),
                }
            ],
        }
        respx.get("https://agent.example.com/.well-known/did.json").respond(json=did_doc_agent)
        respx.get("https://user.example.com/.well-known/did.json").respond(json=did_doc_user)

        result = await verify_receipt(planned_receipt)
        assert result.valid, f"Verification of planned receipt failed: {result.error}"
        print(planned_receipt.to_json())

        result = await verify_receipt(approval_receipt)
        assert result.valid, f"Verification of approval receipt failed: {result.error}"
        print(approval_receipt.to_json())

        result = await verify_receipt(revocation_receipt)
        print(revocation_receipt.to_json())
        for signer in result.signers:
            print(f"Verified signer: {signer.did} with role {signer.role} and key ID {signer.key_id}")
        assert result.valid, f"Verification failed: {result.error}"
