"""Receipt and chain verification."""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

from mailio_receipt.canonicalize import canonicalize, compute_digest
from mailio_receipt.did_resolver import DIDWebResolver
from mailio_receipt.errors import (
    ChainIntegrityError,
    InvalidSignatureError,
    MissingPredecessorError,
)
import base64
from mailio_receipt.signing import verify_jws_detached

if TYPE_CHECKING:
    from mailio_receipt.models import MailioReceipt
    from mailio_receipt.types import DIDResolver


@dataclass(frozen=True)
class SignerInfo:
    """Information about a verified signer."""

    did: str
    role: str
    key_id: str | None


@dataclass(frozen=True)
class VerificationResult:
    """Result of receipt verification."""

    valid: bool
    signers: tuple[SignerInfo, ...]
    error: str | None = None


@dataclass(frozen=True)
class ChainVerificationResult:
    """Result of chain verification."""

    valid: bool
    receipt_results: tuple[VerificationResult, ...]
    chain_intact: bool
    error: str | None = None


async def verify_receipt(
    receipt: MailioReceipt,
    *,
    resolver: DIDResolver | None = None,
) -> VerificationResult:
    """
    Verify a receipt's cryptographic integrity.

    Args:
        receipt: The receipt to verify
        resolver: DID resolver for fetching public keys (default: DIDWebResolver)

    Returns:
        VerificationResult with status and signer details

    Raises:
        InvalidSignatureError: If any signature is invalid
        DIDResolutionError: If signer DID cannot be resolved

    Verification steps:
        1. Recompute eventDigest from event
        2. Verify digest matches receipt.event_digest
        3. For each signature:
           a. Resolve signer DID to get public key
           b. Verify JWS signature over canonicalized event
    """
    if resolver is None:
        resolver = DIDWebResolver()
        owns_resolver = True
    else:
        owns_resolver = False

    try:
        # Step 1 & 2: Verify event digest
        event_dict = receipt.event.to_dict()
        computed_digest = compute_digest(event_dict)

        if computed_digest != receipt.event_digest:
            return VerificationResult(
                valid=False,
                signers=(),
                error=f"Event digest mismatch: expected {receipt.event_digest}, got {computed_digest}",
            )

        # Step 3: Verify all signatures
        canonical_event = canonicalize(event_dict)
        verified_signers: list[SignerInfo] = []

        for sig in receipt.signatures:
            # Resolve signer DID to get public key
            # Extract key ID from JWS header if needed
            key_id = _extract_key_id_from_jws(sig.jws)
            public_key = await resolver.get_verification_key(sig.signer, key_id)

            # Verify JWS signature
            if not verify_jws_detached(sig.jws, canonical_event, public_key):
                raise InvalidSignatureError(f"Invalid signature from {sig.signer}")

            verified_signers.append(
                SignerInfo(
                    did=sig.signer,
                    role=sig.role,
                    key_id=key_id,
                )
            )

        return VerificationResult(
            valid=True,
            signers=tuple(verified_signers),
            error=None,
        )

    except InvalidSignatureError:
        raise
    except Exception as e:
        return VerificationResult(
            valid=False,
            signers=(),
            error=str(e),
        )
    finally:
        if owns_resolver and hasattr(resolver, "close"):
            await resolver.close()


async def verify_chain(
    receipts: list[MailioReceipt],
    *,
    resolver: DIDResolver | None = None,
) -> ChainVerificationResult:
    """
    Verify a chain of receipts for integrity.

    Args:
        receipts: Receipts in chain order (oldest to newest)
        resolver: DID resolver for fetching public keys

    Returns:
        ChainVerificationResult with chain integrity status

    Raises:
        ChainIntegrityError: If chain integrity is violated
        MissingPredecessorError: If a referenced predecessor is not found
        InvalidSignatureError: If any signature is invalid

    Verification steps:
        1. Verify each receipt individually
        2. For each receipt with chain field:
           a. Find predecessor by prevReceiptDigest
           b. Verify predecessor exists in chain
           c. Verify digest matches
    """
    if resolver is None:
        resolver = DIDWebResolver()
        owns_resolver = True
    else:
        owns_resolver = False

    try:
        # Build digest -> receipt mapping
        digest_map: dict[str, MailioReceipt] = {}
        for r in receipts:
            digest_map[r.event_digest] = r

        # Verify each receipt
        receipt_results: list[VerificationResult] = []
        for receipt in receipts:
            result = await verify_receipt(receipt, resolver=resolver)
            receipt_results.append(result)

            if not result.valid:
                return ChainVerificationResult(
                    valid=False,
                    receipt_results=tuple(receipt_results),
                    chain_intact=False,
                    error=result.error,
                )

        # Verify chain integrity
        for receipt in receipts:
            if receipt.event.chain is not None:
                prev_digest = receipt.event.chain.prev_receipt_digest

                # Check if predecessor exists in our chain
                if prev_digest not in digest_map:
                    raise MissingPredecessorError(
                        f"Predecessor {prev_digest[:16]}... not found in chain"
                    )

                # Verify digest matches
                predecessor = digest_map[prev_digest]
                if predecessor.event_digest != prev_digest:
                    raise ChainIntegrityError(
                        f"Chain digest mismatch for predecessor {prev_digest[:16]}..."
                    )

        return ChainVerificationResult(
            valid=True,
            receipt_results=tuple(receipt_results),
            chain_intact=True,
            error=None,
        )

    except (ChainIntegrityError, MissingPredecessorError, InvalidSignatureError):
        raise
    except Exception as e:
        return ChainVerificationResult(
            valid=False,
            receipt_results=tuple(receipt_results) if receipt_results else (),
            chain_intact=False,
            error=str(e),
        )
    finally:
        if owns_resolver and hasattr(resolver, "close"):
            await resolver.close()


def _extract_key_id_from_jws(jws: str) -> str | None:
    """Extract key ID (kid) from JWS header."""
    import base64
    import json

    try:
        header_b64 = jws.split(".")[0]
        # Add padding if needed
        padding = 4 - len(header_b64) % 4
        if padding != 4:
            header_b64 += "=" * padding
        header = json.loads(base64.urlsafe_b64decode(header_b64))
        return header.get("kid")
    except Exception:
        return None
