"""Flask before_request middleware for x402 payment handling."""

from __future__ import annotations

import base64
import json
import os
from datetime import datetime, timezone
from typing import Any, Callable, Optional

try:
    import flask
    from flask import Flask, Request, Response, g, jsonify, request as flask_request
except ImportError as e:
    raise ImportError(
        "flask is required for this module. Install with: pip install x402-stacks[flask]"
    ) from e

from ..types import NetworkType, TokenType, TokenContract
from ..types_v2 import (
    PaymentRequiredV2,
    PaymentRequirementsV2,
    PaymentPayloadV2,
    SettlementResponseV2,
    ResourceInfo,
    X402_HEADERS,
    X402_ERROR_CODES,
)
from ..verifier_v2 import X402PaymentVerifier, SettleOptions
from ..verifier import X402PaymentVerifierV1, SettleOptionsV1
from ..utils import network_to_caip2, asset_to_v2


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


class PaymentMiddlewareConfig:
    """Configuration for Flask x402 payment middleware."""

    def __init__(
        self,
        *,
        network: str,
        amount: str | int,
        pay_to: str,
        scheme: str = "exact",
        asset: Optional[str] = None,
        token_type: Optional[TokenType] = None,
        token_contract: Optional[TokenContract] = None,
        max_timeout_seconds: int = 300,
        facilitator_url: Optional[str] = None,
        description: Optional[str] = None,
        mime_type: Optional[str] = None,
        payment_validator: Optional[Callable[[SettlementResponseV2], bool]] = None,
        extra: Optional[dict[str, Any]] = None,
    ) -> None:
        self.network: str = network if ":" in network else network_to_caip2(network)  # type: ignore[arg-type]
        self.amount = str(amount)
        self.pay_to = pay_to
        self.scheme = scheme
        self.asset = asset or asset_to_v2(token_type or "STX", token_contract)
        self.max_timeout_seconds = max_timeout_seconds
        self.facilitator_url = facilitator_url or "http://localhost:8085"
        self.description = description
        self.mime_type = mime_type
        self.payment_validator = payment_validator
        self.extra = extra


def payment_middleware(config: PaymentMiddlewareConfig) -> Callable[[], Optional[Response]]:
    """Create a Flask before_request function for V2 x402 payments."""
    verifier = X402PaymentVerifier(config.facilitator_url)

    def before_request() -> Optional[Response]:
        sig_header = flask_request.headers.get(X402_HEADERS["PAYMENT_SIGNATURE"])

        if not sig_header:
            return _send_payment_required_v2(config)

        # Decode payment payload
        try:
            decoded = base64.b64decode(sig_header).decode("utf-8")
            payment_payload = PaymentPayloadV2.model_validate(json.loads(decoded))
        except Exception:
            return jsonify({
                "error": X402_ERROR_CODES["INVALID_PAYLOAD"],
                "message": "Invalid payment-signature header",
            }), 400  # type: ignore[return-value]

        if payment_payload.x402_version != 2:
            return jsonify({
                "error": X402_ERROR_CODES["INVALID_X402_VERSION"],
                "message": "Only x402 v2 is supported",
            }), 400  # type: ignore[return-value]

        requirements = PaymentRequirementsV2(
            scheme=config.scheme,
            network=config.network,
            amount=config.amount,
            asset=config.asset,
            pay_to=config.pay_to,
            max_timeout_seconds=config.max_timeout_seconds,
            extra=config.extra,
        )

        settlement = verifier.settle(
            payment_payload, SettleOptions(payment_requirements=requirements)
        )

        if not settlement.success:
            return _send_payment_required_v2(config)

        if config.payment_validator:
            if not config.payment_validator(settlement):
                return jsonify({
                    "error": "custom_validation_failed",
                    "message": "Custom validation rejected the payment",
                }), 402  # type: ignore[return-value]

        g.payment = settlement
        return None

    return before_request


def _send_payment_required_v2(config: PaymentMiddlewareConfig) -> Response:
    """Build and return a 402 response."""
    resource = ResourceInfo(
        url=flask_request.url,
        description=config.description,
        mime_type=config.mime_type,
    )
    requirements = PaymentRequirementsV2(
        scheme=config.scheme,
        network=config.network,
        amount=config.amount,
        asset=config.asset,
        pay_to=config.pay_to,
        max_timeout_seconds=config.max_timeout_seconds,
        extra=config.extra,
    )
    body = PaymentRequiredV2(
        x402_version=2,
        resource=resource,
        accepts=[requirements],
    )
    body_json = body.model_dump(by_alias=True)
    header_val = base64.b64encode(json.dumps(body_json).encode()).decode()

    resp = jsonify(body_json)
    resp.status_code = 402
    resp.headers[X402_HEADERS["PAYMENT_REQUIRED"]] = header_val
    return resp


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


def payment_middleware_v1(
    *,
    amount: str | int,
    address: str,
    network: NetworkType,
    facilitator_url: str = "http://localhost:8085",
    resource: Optional[str] = None,
    expiration_seconds: int = 300,
    token_type: Optional[TokenType] = None,
    token_contract: Optional[TokenContract] = None,
) -> Callable[[], Optional[Response]]:
    """Create a Flask before_request function for V1 x402 payments."""
    verifier = X402PaymentVerifierV1(facilitator_url, network)

    def before_request() -> Optional[Response]:
        signed_payment = flask_request.headers.get("X-PAYMENT")

        if not signed_payment:
            return _send_payment_required_v1(
                amount=amount,
                address=address,
                network=network,
                resource=resource,
                expiration_seconds=expiration_seconds,
                token_type=token_type,
                token_contract=token_contract,
            )

        header_token = (flask_request.headers.get("X-PAYMENT-TOKEN-TYPE") or "").lower()
        tt: TokenType = token_type or "STX"
        if header_token == "sbtc":
            tt = "sBTC"
        elif header_token == "usdcx":
            tt = "USDCx"

        settle_opts = SettleOptionsV1(
            expected_recipient=address,
            min_amount=int(amount),
            token_type=tt,
            resource=resource or flask_request.path,
            method=flask_request.method,
        )
        verification = verifier.settle_payment(signed_payment, settle_opts)

        if not verification.is_valid:
            return jsonify({
                "error": "Invalid payment",
                "details": verification.validation_error,
                "paymentStatus": verification.status,
            }), 402  # type: ignore[return-value]

        if verification.status == "pending":
            return jsonify({
                "error": "Payment not yet confirmed",
                "paymentStatus": "pending",
            }), 402  # type: ignore[return-value]

        if verification.status == "failed":
            return jsonify({
                "error": "Payment failed",
                "paymentStatus": "failed",
            }), 402  # type: ignore[return-value]

        g.payment = verification
        return None

    return before_request


def _send_payment_required_v1(
    *,
    amount: str | int,
    address: str,
    network: NetworkType,
    resource: Optional[str],
    expiration_seconds: int,
    token_type: Optional[TokenType],
    token_contract: Optional[TokenContract],
) -> Response:
    nonce = os.urandom(16).hex()
    expires_at = datetime.fromtimestamp(
        datetime.now(timezone.utc).timestamp() + expiration_seconds,
        tz=timezone.utc,
    ).isoformat().replace("+00:00", "Z")

    body: dict[str, Any] = {
        "maxAmountRequired": str(amount),
        "resource": resource or flask_request.path,
        "payTo": address,
        "network": network,
        "nonce": nonce,
        "expiresAt": expires_at,
    }
    if token_type:
        body["tokenType"] = token_type
    if token_contract:
        body["tokenContract"] = token_contract.model_dump()

    resp = jsonify(body)
    resp.status_code = 402
    return resp


# ── Helpers ──────────────────────────────────────────────────────────


def get_payment() -> Optional[SettlementResponseV2]:
    """Get payment info from Flask g context."""
    return getattr(g, "payment", None)


def conditional_payment(
    condition: Callable[[], bool],
    config: PaymentMiddlewareConfig,
) -> Callable[[], Optional[Response]]:
    """Only require payment if condition is met."""
    mw = payment_middleware(config)

    def before_request() -> Optional[Response]:
        if condition():
            return mw()
        return None

    return before_request


def tiered_payment(
    get_tier: Callable[[], dict[str, Any]],
    base_config: dict[str, Any],
) -> Callable[[], Optional[Response]]:
    """Different amounts based on request."""

    def before_request() -> Optional[Response]:
        tier = get_tier()
        cfg = PaymentMiddlewareConfig(**{**base_config, **tier})
        mw = payment_middleware(cfg)
        return mw()

    return before_request


def payment_rate_limit(
    *,
    free_requests: int,
    window_ms: int,
    payment_config: PaymentMiddlewareConfig,
    key_generator: Optional[Callable[[], str]] = None,
) -> Callable[[], Optional[Response]]:
    """Rate limiting with payment after free tier."""
    request_counts: dict[str, dict[str, Any]] = {}
    mw = payment_middleware(payment_config)

    def before_request() -> Optional[Response]:
        key = key_generator() if key_generator else (flask_request.remote_addr or "unknown")
        now = int(datetime.now(timezone.utc).timestamp() * 1000)

        record = request_counts.get(key)
        if record and record["reset_at"] < now:
            record = None

        if not record:
            record = {"count": 0, "reset_at": now + window_ms}
            request_counts[key] = record

        if record["count"] >= free_requests:
            return mw()

        record["count"] += 1
        return None

    return before_request
