"""httpx transport wrappers for automatic x402 payment handling."""

from __future__ import annotations

from typing import Any, Optional

try:
    import httpx
except ImportError as e:
    raise ImportError(
        "httpx is required for this module. Install with: pip install x402-stacks[httpx]"
    ) from e

from ..types import StacksAccount, X402PaymentRequired
from ..types_v2 import (
    PaymentPayloadV2,
    PaymentRequiredV2,
    StacksPayloadV2,
    X402_HEADERS,
)
from .interceptor import (
    sign_payment_for_request,
    is_valid_payment_request,
    decode_x_payment_response,
    private_key_to_account_v1,
)
from .interceptor_v2 import (
    sign_payment_v2,
    select_payment_option,
    decode_payment_required,
    encode_payment_payload,
    decode_payment_response,
    private_key_to_account,
)


# ── V2 (default) ────────────────────────────────────────────────────


class _PaymentTransport(httpx.AsyncBaseTransport):
    """Async transport that intercepts 402 responses and retries with payment."""

    def __init__(
        self,
        transport: httpx.AsyncBaseTransport,
        account: StacksAccount,
    ) -> None:
        self._transport = transport
        self._account = account

    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
        response = await self._transport.handle_async_request(request)

        if response.status_code != 402:
            return response

        # Already attempted?
        if request.headers.get(X402_HEADERS["PAYMENT_SIGNATURE"]):
            return response

        # Try header first, then body
        await response.aread()
        payment_required: Optional[PaymentRequiredV2] = None

        header_val = response.headers.get(X402_HEADERS["PAYMENT_REQUIRED"])
        if header_val:
            payment_required = decode_payment_required(header_val)

        if not payment_required:
            try:
                body = response.json()
                if (
                    isinstance(body, dict)
                    and body.get("x402Version") == 2
                    and isinstance(body.get("accepts"), list)
                ):
                    payment_required = PaymentRequiredV2.model_validate(body)
            except Exception:
                pass

        if not payment_required:
            return response

        selected = select_payment_option(payment_required.accepts, self._account)
        if not selected:
            return response

        signed_tx = sign_payment_v2(selected, self._account)
        payload = PaymentPayloadV2(
            x402_version=2,
            resource=payment_required.resource,
            accepted=selected,
            payload=StacksPayloadV2(transaction=signed_tx),
        )
        encoded = encode_payment_payload(payload)

        retry_request = httpx.Request(
            method=request.method,
            url=request.url,
            headers={**dict(request.headers), X402_HEADERS["PAYMENT_SIGNATURE"]: encoded},
            content=request.content,
        )
        return await self._transport.handle_async_request(retry_request)


class _PaymentTransportSync(httpx.BaseTransport):
    """Sync transport that intercepts 402 responses and retries with payment."""

    def __init__(
        self,
        transport: httpx.BaseTransport,
        account: StacksAccount,
    ) -> None:
        self._transport = transport
        self._account = account

    def handle_request(self, request: httpx.Request) -> httpx.Response:
        response = self._transport.handle_request(request)

        if response.status_code != 402:
            return response

        if request.headers.get(X402_HEADERS["PAYMENT_SIGNATURE"]):
            return response

        response.read()
        payment_required: Optional[PaymentRequiredV2] = None

        header_val = response.headers.get(X402_HEADERS["PAYMENT_REQUIRED"])
        if header_val:
            payment_required = decode_payment_required(header_val)

        if not payment_required:
            try:
                body = response.json()
                if (
                    isinstance(body, dict)
                    and body.get("x402Version") == 2
                    and isinstance(body.get("accepts"), list)
                ):
                    payment_required = PaymentRequiredV2.model_validate(body)
            except Exception:
                pass

        if not payment_required:
            return response

        selected = select_payment_option(payment_required.accepts, self._account)
        if not selected:
            return response

        signed_tx = sign_payment_v2(selected, self._account)
        payload = PaymentPayloadV2(
            x402_version=2,
            resource=payment_required.resource,
            accepted=selected,
            payload=StacksPayloadV2(transaction=signed_tx),
        )
        encoded = encode_payment_payload(payload)

        retry_request = httpx.Request(
            method=request.method,
            url=request.url,
            headers={**dict(request.headers), X402_HEADERS["PAYMENT_SIGNATURE"]: encoded},
            content=request.content,
        )
        return self._transport.handle_request(retry_request)


def wrap_httpx_with_payment(
    client: httpx.AsyncClient, account: StacksAccount
) -> httpx.AsyncClient:
    """Wrap an async httpx client with V2 payment handling."""
    client._transport = _PaymentTransport(client._transport, account)  # type: ignore[assignment]
    return client


def wrap_httpx_with_payment_sync(
    client: httpx.Client, account: StacksAccount
) -> httpx.Client:
    """Wrap a sync httpx client with V2 payment handling."""
    client._transport = _PaymentTransportSync(client._transport, account)  # type: ignore[assignment]
    return client


def create_payment_client(account: StacksAccount, **kwargs: Any) -> httpx.AsyncClient:
    """Create an async httpx client with V2 payment handling."""
    client = httpx.AsyncClient(**kwargs)
    return wrap_httpx_with_payment(client, account)


def create_payment_client_sync(account: StacksAccount, **kwargs: Any) -> httpx.Client:
    """Create a sync httpx client with V2 payment handling."""
    client = httpx.Client(**kwargs)
    return wrap_httpx_with_payment_sync(client, account)


# ── V1 (legacy) ─────────────────────────────────────────────────────


class _PaymentTransportV1(httpx.AsyncBaseTransport):
    """Async transport for V1 x402 protocol."""

    def __init__(
        self,
        transport: httpx.AsyncBaseTransport,
        account: StacksAccount,
    ) -> None:
        self._transport = transport
        self._account = account

    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
        response = await self._transport.handle_async_request(request)

        if response.status_code != 402:
            return response

        if request.headers.get("x-payment"):
            return response

        await response.aread()
        try:
            body = response.json()
        except Exception:
            return response

        if not is_valid_payment_request(body):
            return response

        payment_req = X402PaymentRequired.model_validate(body)
        signed_tx = sign_payment_for_request(payment_req, self._account)

        retry_request = httpx.Request(
            method=request.method,
            url=request.url,
            headers={
                **dict(request.headers),
                "x-payment": signed_tx,
                "x-payment-token-type": payment_req.token_type or "STX",
            },
            content=request.content,
        )
        return await self._transport.handle_async_request(retry_request)


class _PaymentTransportV1Sync(httpx.BaseTransport):
    """Sync transport for V1 x402 protocol."""

    def __init__(
        self,
        transport: httpx.BaseTransport,
        account: StacksAccount,
    ) -> None:
        self._transport = transport
        self._account = account

    def handle_request(self, request: httpx.Request) -> httpx.Response:
        response = self._transport.handle_request(request)

        if response.status_code != 402:
            return response

        if request.headers.get("x-payment"):
            return response

        response.read()
        try:
            body = response.json()
        except Exception:
            return response

        if not is_valid_payment_request(body):
            return response

        payment_req = X402PaymentRequired.model_validate(body)
        signed_tx = sign_payment_for_request(payment_req, self._account)

        retry_request = httpx.Request(
            method=request.method,
            url=request.url,
            headers={
                **dict(request.headers),
                "x-payment": signed_tx,
                "x-payment-token-type": payment_req.token_type or "STX",
            },
            content=request.content,
        )
        return self._transport.handle_request(retry_request)


def wrap_httpx_with_payment_v1(
    client: httpx.AsyncClient, account: StacksAccount
) -> httpx.AsyncClient:
    """Wrap an async httpx client with V1 payment handling."""
    client._transport = _PaymentTransportV1(client._transport, account)  # type: ignore[assignment]
    return client


def wrap_httpx_with_payment_v1_sync(
    client: httpx.Client, account: StacksAccount
) -> httpx.Client:
    """Wrap a sync httpx client with V1 payment handling."""
    client._transport = _PaymentTransportV1Sync(client._transport, account)  # type: ignore[assignment]
    return client


def create_payment_client_v1(account: StacksAccount, **kwargs: Any) -> httpx.AsyncClient:
    """Create an async httpx client with V1 payment handling."""
    client = httpx.AsyncClient(**kwargs)
    return wrap_httpx_with_payment_v1(client, account)


def create_payment_client_v1_sync(account: StacksAccount, **kwargs: Any) -> httpx.Client:
    """Create a sync httpx client with V1 payment handling."""
    client = httpx.Client(**kwargs)
    return wrap_httpx_with_payment_v1_sync(client, account)
