"""Receipt and chain verification."""

from __future__ import annotations

import asyncio
import base64
import json
import urllib.parse
from dataclasses import dataclass
from typing import TYPE_CHECKING

import cbor2
import httpx
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.ec import (
    SECP256R1,
    EllipticCurvePublicNumbers,
)
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature

from mailio_receipt.canonicalize import canonicalize, compute_digest
from mailio_receipt.errors import (
    ChainIntegrityError,
    InvalidSignatureError,
    MissingPredecessorError,
    SignerNotFoundError,
)
from mailio_receipt.signing import verify_jws_detached

if TYPE_CHECKING:
    from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurvePublicKey

    from mailio_receipt.models import MailioReceipt

    PublicKey = EllipticCurvePublicKey | Ed25519PublicKey | WebAuthnPublicKey


@dataclass(frozen=True)
class WebAuthnPublicKey:
    """Wrapper to distinguish WebAuthn COSE keys from DID JWK keys."""

    key: EllipticCurvePublicKey


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

    signer: 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

def _parse_cose_key(b64_public_key: str) -> EllipticCurvePublicKey:
    """Parse a base64-encoded COSE public key (P-256/ES256) into a cryptography key."""
    raw = base64.b64decode(b64_public_key)
    cose = cbor2.loads(raw)

    # COSE key labels: 1=kty, 3=alg, -1=crv, -2=x, -3=y
    kty = cose.get(1)
    crv = cose.get(-1)

    if kty != 2:  # EC2
        raise ValueError(f"Unsupported COSE key type: {kty}")
    if crv != 1:  # P-256
        raise ValueError(f"Unsupported COSE curve: {crv}")

    x = int.from_bytes(cose[-2], "big")
    y = int.from_bytes(cose[-3], "big")
    numbers = EllipticCurvePublicNumbers(x=x, y=y, curve=SECP256R1())
    return numbers.public_key()


def _b64url_decode(s: str) -> bytes:
    """Base64url decode with padding restoration."""
    padding = 4 - len(s) % 4
    if padding != 4:
        s += "=" * padding
    return base64.urlsafe_b64decode(s)


def _decode_signature(signature_b64: str) -> bytes:
    """Decode a base64url-encoded ECDSA signature (DER or R||S) to DER."""
    signature_raw = _b64url_decode(signature_b64)

    if len(signature_raw) == 64:
        r = int.from_bytes(signature_raw[:32], byteorder="big")
        s = int.from_bytes(signature_raw[32:], byteorder="big")
        return encode_dss_signature(r, s)
    return signature_raw


def _verify_webauthn_signature(
    signature_b64: str,
    public_key: EllipticCurvePublicKey,
    authenticator_data_b64: str,
    client_data_json_b64: str,
) -> bool:
    """Verify a WebAuthn assertion signature.

    The WebAuthn authenticator signs ``authenticatorData || SHA-256(clientDataJSON)``.
    This function reconstructs that signed payload and verifies the ECDSA signature.
    """
    try:
        import hashlib

        signature_der = _decode_signature(signature_b64)
        authenticator_data = _b64url_decode(authenticator_data_b64)
        client_data_json_bytes = _b64url_decode(client_data_json_b64)

        client_data_hash = hashlib.sha256(client_data_json_bytes).digest()
        signed_data = authenticator_data + client_data_hash

        public_key.verify(signature_der, signed_data, ec.ECDSA(hashes.SHA256()))
        return True
    except InvalidSignature:
        return False
    except Exception:
        return False


def _b64url_decode_jwk(s: str) -> bytes:
    """Decode a base64url-encoded JWK value (with padding restoration)."""
    padding = 4 - len(s) % 4
    if padding != 4:
        s += "=" * padding
    return base64.urlsafe_b64decode(s)


def _parse_jwk(jwk: dict) -> PublicKey:
    """Parse a JWK dict into a cryptography public key."""
    kty = jwk.get("kty")
    if kty == "EC":
        crv = jwk.get("crv")
        if crv != "P-256":
            raise ValueError(f"Unsupported EC curve: {crv}")
        x = int.from_bytes(_b64url_decode_jwk(jwk["x"]), "big")
        y = int.from_bytes(_b64url_decode_jwk(jwk["y"]), "big")
        numbers = EllipticCurvePublicNumbers(x=x, y=y, curve=SECP256R1())
        return numbers.public_key()
    elif kty == "OKP":
        crv = jwk.get("crv")
        if crv != "Ed25519":
            raise ValueError(f"Unsupported OKP curve: {crv}")
        x_bytes = _b64url_decode_jwk(jwk["x"])
        return Ed25519PublicKey.from_public_bytes(x_bytes)
    else:
        raise ValueError(f"Unsupported JWK key type: {kty}")


def _did_to_url(signer: str) -> str:
    """Resolve a did:web identifier to its DID document URL.

    did:web:mio.mailiomail.com:0xabc#master -> https://mio.mailiomail.com/0xabc/did.json
    """
    # Strip fragment (e.g. #master) before parsing
    did_part = signer.split("#")[0]
    parts = did_part.split(":")
    # parts = ["did", "web", domain, ...path_segments]
    domain = parts[2].replace("%3A", ":")
    path = "/".join(parts[3:])
    return f"https://{domain}/{path}/did.json"


def _parse_did_document(did_doc: dict, signer: str) -> PublicKey:
    """Extract a public key from a DID document's verificationMethod matching the signer.

    The signer includes a fragment (e.g. did:web:domain:0xabc#master) that identifies
    which verificationMethod entry to use.
    """
    methods = did_doc.get("verificationMethod", [])
    if not methods:
        raise ValueError("DID document has no verificationMethod entries")

    fragment = signer.split("#")[1] if "#" in signer else None

    for method in methods:
        method_id = method.get("id", "")
        # Match by full signer, or by fragment (e.g. #master)
        if method_id == signer or (fragment and method_id.endswith(f"#{fragment}")):
            if "publicKeyJwk" in method:
                return _parse_jwk(method["publicKeyJwk"])
            raise ValueError(f"verificationMethod '{method_id}' has no publicKeyJwk")

    raise ValueError(f"No verificationMethod matching '{signer}' found in DID document")


async def _fetch_did_key(
    client: httpx.AsyncClient, signer: str
) -> tuple[str, PublicKey]:
    """Fetch and parse a public key from a DID document."""
    url = _did_to_url(signer)
    resp = await client.get(url)
    resp.raise_for_status()
    did_doc = resp.json()
    return (signer, _parse_did_document(did_doc, signer))


async def _fetch_webauthn_key(
    client: httpx.AsyncClient, signer: str, mio_backend_url: str, api_key: str
) -> tuple[str, PublicKey]:
    """Fetch and parse a COSE-formatted public key from the webauthn API endpoint."""
    encoded_signer = urllib.parse.quote(signer, safe="")
    url = f"https://{mio_backend_url}/api/v1/receipts/webauthn/public_key/{encoded_signer}"
    headers = {"X-Mailio-Service-Key": api_key}
    resp = await client.get(url, headers=headers)
    resp.raise_for_status()
    key_data = resp.json()
    return (signer, WebAuthnPublicKey(key=_parse_cose_key(key_data["publicKey"])))


async def query_backend_public_keys(
    receipt: MailioReceipt,
    mio_backend_url: str = "mio.mailiomail.com",
    api_key: str = "",
) -> dict[str, PublicKey]:
    """
    Mio has 2 types of public keys (did public keys used by agents and webauthn public keys belonging to users)
    1. DID are public keys retrievable by pattern: https://mio.servername.com/0xabc/did.json
    2. Webauthn are users public keys retrieved by endpoint call to
       https://mio.servername/api/v1/receipts/publickey/{signer}
    Return:
        dict[signer:publicKey]
    """
    if receipt.signatures is None or len(receipt.signatures) == 0:
        return {}

    tasks: list[asyncio.Task[tuple[str, PublicKey]]] = []
    async with httpx.AsyncClient() as client:
        for signature in receipt.signatures:
            signer = signature.signer
            if signer is None:
                continue
            if "did" in signer:
                tasks.append(asyncio.ensure_future(_fetch_did_key(client, signer)))
            elif "webauthn" in signer:
                tasks.append(
                    asyncio.ensure_future(
                        _fetch_webauthn_key(client, signer, mio_backend_url, api_key)
                    )
                )
            else:
                raise ValueError(f"Unknown key type for signer: {signer}")

        results = await asyncio.gather(*tasks)

    return dict(results)

async def verify_receipt(
    receipt: MailioReceipt,
    *,
    mio_backend_url: str,
    mio_backend_api_key: str,
) -> VerificationResult:
    """
    Verify a receipt's cryptographic integrity.

    Args:
        receipt: The receipt to verify
        public_keys: Mapping of signer identifier to public key.
            The caller is responsible for obtaining keys (via DID resolution,
            WebAuthn, database lookup, etc.).

    Returns:
        VerificationResult with statusq
         and signer details

    Raises:
        InvalidSignatureError: If any signature is invalid
        SignerNotFoundError: If a signer's public key is not in public_keys

    Verification steps:
        1. Recompute eventDigest from event
        2. Verify digest matches receipt.event_digest
        3. For each signature:
           a. Look up signer's public key from public_keys
           b. Verify JWS signature over canonicalized event
    """
    # 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: Query backend API keys
    public_keys = await query_backend_public_keys(receipt, mio_backend_url, mio_backend_api_key)

    verified_signers: list[SignerInfo] = []

    for sig in receipt.signatures:
        public_key = public_keys.get(sig.signer)
        if public_key is None:
            raise SignerNotFoundError(
                f"No public key provided for signer: {sig.signer}"
            )

        if isinstance(public_key, WebAuthnPublicKey):
            key_id = None
            if sig.authenticator_data is None or sig.client_data_json is None:
                raise InvalidSignatureError(
                    f"WebAuthn signature from {sig.signer} missing authenticator_data or client_data_json"
                )
            if not _verify_webauthn_signature(
                sig.jws, public_key.key, sig.authenticator_data, sig.client_data_json
            ):
                raise InvalidSignatureError(f"Invalid signature from {sig.signer}")
        else:
            key_id = _extract_key_id_from_jws(sig.jws)
            if not verify_jws_detached(sig.jws, computed_digest, public_key):
                raise InvalidSignatureError(f"Invalid signature from {sig.signer}")

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

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


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

    Args:
        receipts: Receipts in chain order (oldest to newest)
        public_keys: Mapping of signer identifier to public key

    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
        SignerNotFoundError: If a signer's public key is not in public_keys

    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
    """
    receipt_results: list[VerificationResult] = []

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

        # Verify each receipt
        for receipt in receipts:
            result = await verify_receipt(receipt, public_keys=public_keys)
            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

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

                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, SignerNotFoundError):
        raise
    except Exception as e:
        return ChainVerificationResult(
            valid=False,
            receipt_results=tuple(receipt_results) if receipt_results else (),
            chain_intact=False,
            error=str(e),
        )


def _extract_key_id_from_jws(jws: str) -> str | None:
    """Extract key ID (kid) from JWS header."""
    try:
        header_b64 = jws.split(".")[0]
        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
