"""requests.Session subclass with automatic x402 payment handling."""

from __future__ import annotations

from typing import Any, Optional

try:
    import requests
except ImportError as e:
    raise ImportError(
        "requests is required for this module. Install with: pip install x402-stacks[requests]"
    ) 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,
)
from .interceptor_v2 import (
    sign_payment_v2,
    select_payment_option,
    decode_payment_required,
    encode_payment_payload,
)


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


class PaymentSession(requests.Session):
    """requests.Session that automatically handles 402 payment responses (V2)."""

    def __init__(self, account: StacksAccount, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._account = account

    def request(self, method: str, url: str, **kwargs: Any) -> requests.Response:  # type: ignore[override]
        response = super().request(method, url, **kwargs)

        if response.status_code != 402:
            return response

        # Already attempted?
        headers = kwargs.get("headers", {}) or {}
        if X402_HEADERS["PAYMENT_SIGNATURE"] in headers:
            return response

        # Try header first, then body
        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 with payment
        retry_headers = dict(headers)
        retry_headers[X402_HEADERS["PAYMENT_SIGNATURE"]] = encoded
        kwargs["headers"] = retry_headers
        return super().request(method, url, **kwargs)


def create_payment_session(account: StacksAccount) -> PaymentSession:
    """Create a requests session with V2 payment handling."""
    return PaymentSession(account)


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


class PaymentSessionV1(requests.Session):
    """requests.Session that automatically handles 402 payment responses (V1)."""

    def __init__(self, account: StacksAccount, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._account = account

    def request(self, method: str, url: str, **kwargs: Any) -> requests.Response:  # type: ignore[override]
        response = super().request(method, url, **kwargs)

        if response.status_code != 402:
            return response

        headers = kwargs.get("headers", {}) or {}
        if "x-payment" in headers or "X-PAYMENT" in headers:
            return response

        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_headers = dict(headers)
        retry_headers["X-PAYMENT"] = signed_tx
        retry_headers["X-PAYMENT-TOKEN-TYPE"] = payment_req.token_type or "STX"
        kwargs["headers"] = retry_headers
        return super().request(method, url, **kwargs)


def create_payment_session_v1(account: StacksAccount) -> PaymentSessionV1:
    """Create a requests session with V1 payment handling."""
    return PaymentSessionV1(account)
