"""Unified MailioClient for creating, signing, verifying, and fetching receipts."""

from __future__ import annotations

import asyncio
import base64
import logging
import random
from dataclasses import dataclass
from typing import Any, Literal

import httpx
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

from mailio_receipt.canonicalize import canonicalize, compute_digest
from mailio_receipt.errors import SignatureError, ValidationError
from mailio_receipt.models import (
    MailioAuditEvent,
    MailioReceipt,
    Signature,
    SignedMetadata,
)
from mailio_receipt.signing import create_jws_detached
from mailio_receipt.verification import VerificationResult
from mailio_receipt.verification import verify_receipt as _verify_receipt

logger = logging.getLogger(__name__)


def _parse_retry_after(value: str | None) -> float | None:
    """Parse a Retry-After header value (seconds) or return None."""
    if value is None:
        return None
    try:
        return float(value)
    except ValueError:
        return None


@dataclass(frozen=True)
class CreateReceiptResult:
    """Result of MailioClient.create_receipt().

    Always contains a valid receipt, even if saving to the backend failed.

    Attributes:
        receipt: The signed MailioReceipt.
        saved: True if the receipt was successfully persisted to the backend.
        save_error: The exception that caused the save to fail, or None.
    """

    receipt: MailioReceipt
    saved: bool
    save_error: Exception | None = None


class MailioClient:
    """High-level client for creating, signing, verifying, and fetching receipts.

    Accepts the agent config dict (provided at registration) which contains
    all keys, DID, and API key. All methods use these stored credentials
    so callers only need to provide business parameters.

    Usage::

        async with MailioClient(config) as client:
            result = await client.create_receipt(event)
            receipt = await client.get_receipt(receipt_id)
    """

    def __init__(
        self,
        config: dict[str, Any],
        *,
        max_retries: int = 3,
        timeout: float = 30.0,
        backoff_base: float = 0.5,
        backoff_max: float = 30.0,
    ) -> None:
        # Extract private key (base64 → 64 bytes, first 32 = seed)
        key_bytes = base64.b64decode(config["signingPrivateKey"])
        self._private_key = Ed25519PrivateKey.from_private_bytes(key_bytes[:32])

        # Extract signer DID
        self._signer_did: str = config["did"]

        # Extract backend host from didDocumentStoredKey
        # Format: did:web:<hostname>:<address>
        did_parts = config["didDocumentStoredKey"].split(":")
        self._backend_host: str = did_parts[2]
        self._base_url: str = f"https://{self._backend_host}"

        # API key
        self._api_key: str = config["api_key"]

        # Retry configuration
        self._max_retries = max_retries
        self._backoff_base = backoff_base
        self._backoff_max = backoff_max

        # Shared HTTP client
        self._client = httpx.AsyncClient(
            base_url=self._base_url,
            headers={"X-API-Key": self._api_key},
            timeout=httpx.Timeout(timeout),
        )

    async def __aenter__(self) -> MailioClient:
        return self

    async def __aexit__(self, *exc: object) -> None:
        await self.close()

    async def close(self) -> None:
        """Close the underlying HTTP client."""
        await self._client.aclose()

    # ── HTTP resilience ──────────────────────────────────────────────

    def _compute_delay(self, attempt: int) -> float:
        """Exponential backoff with ±25% jitter, capped at backoff_max."""
        delay = self._backoff_base * (2 ** attempt)
        delay = min(delay, self._backoff_max)
        jitter = delay * 0.25 * (2 * random.random() - 1)  # noqa: S311
        return max(0.0, delay + jitter)

    async def _request(
        self,
        method: str,
        path: str,
        *,
        json: dict[str, Any] | None = None,
    ) -> httpx.Response:
        """Send an HTTP request with retry and exponential backoff.

        Retries on 429 (rate-limit) and 5xx (server error) responses,
        as well as on timeouts. Raises immediately on 4xx (client error).
        """
        last_exc: Exception | None = None

        for attempt in range(self._max_retries + 1):
            try:
                resp = await self._client.request(method, path, json=json)

                # 2xx — success
                if resp.is_success:
                    return resp

                # 4xx (not 429) — client error, don't retry
                if 400 <= resp.status_code < 500 and resp.status_code != 429:
                    resp.raise_for_status()

                # 429 — rate limited
                if resp.status_code == 429:
                    if attempt < self._max_retries:
                        retry_after = _parse_retry_after(
                            resp.headers.get("Retry-After")
                        )
                        delay = retry_after if retry_after is not None else self._compute_delay(attempt)
                        logger.warning(
                            "Rate limited (429), retrying in %.1fs (attempt %d/%d)",
                            delay, attempt + 1, self._max_retries,
                        )
                        await asyncio.sleep(delay)
                        continue
                    resp.raise_for_status()

                # 5xx — server error
                if resp.status_code >= 500:
                    if attempt < self._max_retries:
                        delay = self._compute_delay(attempt)
                        logger.warning(
                            "Server error %d, retrying in %.1fs (attempt %d/%d)",
                            resp.status_code, delay, attempt + 1, self._max_retries,
                        )
                        await asyncio.sleep(delay)
                        continue
                    resp.raise_for_status()

            except httpx.TimeoutException as exc:
                last_exc = exc
                if attempt < self._max_retries:
                    delay = self._compute_delay(attempt)
                    logger.warning(
                        "Timeout, retrying in %.1fs (attempt %d/%d)",
                        delay, attempt + 1, self._max_retries,
                    )
                    await asyncio.sleep(delay)
                    continue
                raise
            except httpx.HTTPStatusError:
                raise

        # Safety net — should not normally reach here
        if last_exc is not None:
            raise last_exc
        msg = "Request failed after all retries"
        raise httpx.TransportError(msg)

    # ── Public API ───────────────────────────────────────────────────

    async def create_receipt(
        self,
        event: MailioAuditEvent,
        role: Literal["user", "agent", "authority"] = "agent",
        *,
        save: bool = True,
    ) -> CreateReceiptResult:
        """Create a signed receipt and optionally persist it to the backend.

        The receipt is always returned, even if saving fails after retries.

        Args:
            event: The event to receipt (must be fully populated).
            role: Role of the signer in this event.
            save: If True, POST the receipt to the backend with retries.

        Returns:
            CreateReceiptResult with the receipt and save status.
        """
        # Sign locally
        event_dict = event.to_dict()
        event_digest = compute_digest(event_dict)
        canonical_event = canonicalize(event_dict)

        try:
            jws = create_jws_detached(
                canonical_event, self._private_key, self._signer_did
            )
        except Exception as e:
            raise SignatureError(f"Failed to create signature: {e}") from e

        signed_metadata = SignedMetadata(
            canonicalization="RFC8785",
            digest=event_digest,
            target="/event",
        )
        signature = Signature(
            role=role,
            signer=self._signer_did,
            alg="EdDSA",
            signed=signed_metadata,
            jws=jws,
        )
        receipt = MailioReceipt(
            type="MailioReceipt",
            v="v1",
            event=event,
            event_digest=event_digest,
            signatures=(signature,),
        )

        if not save:
            return CreateReceiptResult(receipt=receipt, saved=False)

        # Save with retries
        try:
            resp = await self._request(
                "POST", "/api/v1/receipts", json=receipt.to_dict()
            )
            saved_receipt = MailioReceipt.from_dict(resp.json())
            return CreateReceiptResult(receipt=saved_receipt, saved=True)
        except Exception as exc:
            return CreateReceiptResult(
                receipt=receipt, saved=False, save_error=exc
            )

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

        Uses the backend configured in the client for public key lookup.
        """
        return await _verify_receipt(
            receipt,
            mio_backend_url=self._backend_host,
            mio_backend_api_key=self._api_key,
        )

    async def sign_receipt(
        self,
        receipt: MailioReceipt,
        role: Literal["user", "agent", "authority"] = "agent",
        *,
        save: bool = True,
    ) -> CreateReceiptResult:
        """Add a co-signature to an existing receipt.

        Returns a CreateReceiptResult containing a new immutable MailioReceipt
        with the additional signature appended. The original receipt is not
        modified. If save=True, the co-signed receipt is persisted to the
        backend with retries.

        Args:
            receipt: The receipt to co-sign.
            role: Role of the co-signer.
            save: If True, POST the co-signed receipt to the backend.

        Returns:
            CreateReceiptResult with the co-signed receipt and save status.
        """
        event_dict = receipt.event.to_dict()
        canonical_event = canonicalize(event_dict)

        try:
            jws = create_jws_detached(
                canonical_event, self._private_key, self._signer_did
            )
        except Exception as e:
            raise SignatureError(f"Failed to create co-signature: {e}") from e

        signed_metadata = SignedMetadata(
            canonicalization="RFC8785",
            digest=receipt.event_digest,
            target="/event",
        )
        new_sig = Signature(
            role=role,
            signer=self._signer_did,
            alg="EdDSA",
            signed=signed_metadata,
            jws=jws,
        )

        cosigned = MailioReceipt(
            type=receipt.type,
            v=receipt.v,
            event=receipt.event,
            event_digest=receipt.event_digest,
            signatures=receipt.signatures + (new_sig,),
        )

        if not save:
            return CreateReceiptResult(receipt=cosigned, saved=False)

        try:
            resp = await self._request(
                "POST", "/api/v1/receipts", json=cosigned.to_dict()
            )
            saved_receipt = MailioReceipt.from_dict(resp.json())
            return CreateReceiptResult(receipt=saved_receipt, saved=True)
        except Exception as exc:
            return CreateReceiptResult(
                receipt=cosigned, saved=False, save_error=exc
            )

    async def save_receipt(self, receipt: MailioReceipt) -> CreateReceiptResult:
        """Save an existing receipt to the backend.

        Useful when you have a receipt created with save=False (or from
        another source) and want to persist it separately.

        Args:
            receipt: The receipt to save.

        Returns:
            CreateReceiptResult with the save status.
        """
        try:
            resp = await self._request(
                "POST", "/api/v1/receipts", json=receipt.to_dict()
            )
            saved_receipt = MailioReceipt.from_dict(resp.json())
            return CreateReceiptResult(receipt=saved_receipt, saved=True)
        except Exception as exc:
            return CreateReceiptResult(
                receipt=receipt, saved=False, save_error=exc
            )

    async def get_receipt(self, receipt_id: str) -> MailioReceipt:
        """Fetch a single receipt by ID from the backend."""
        resp = await self._request("GET", f"/api/v1/receipts/{receipt_id}")
        return MailioReceipt.from_dict(resp.json())

    async def get_flow_receipts(self, flow_id: str) -> list[MailioReceipt]:
        """Fetch all receipts belonging to a flow."""
        resp = await self._request(
            "GET", f"/api/v1/receipts/flow/{flow_id}"
        )
        data = resp.json()
        return [MailioReceipt.from_dict(r) for r in data.get("receipts", [])]

    async def mint_registration_token(self, external_user_id: str) -> dict[str, Any]:
        """Mint a WebAuthn registration token via the Mailio API.

        Args:
            external_user_id: The external user identifier.

        Returns:
            The registration token response from the backend.

        Raises:
            ValidationError: If external_user_id is empty.
            httpx.HTTPStatusError: If the backend returns an error.
        """
        if not external_user_id or not external_user_id.strip():
            raise ValidationError("external_user_id must not be empty")

        resp = await self._request(
            "POST",
            "/api/v1/receipts/webauthn/mint_registration_token",
            json={"externalUserId": external_user_id},
        )
        return resp.json()
