routstr-core/routstr/balance.py

import hashlib
from typing import Annotated, NoReturn

from fastapi import APIRouter, Depends, Header, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from sqlmodel import col, select

from . import refund
from .auth import (
    redemption_error_to_http_exception,
    validate_bearer_key,
)
from .core.db import (
    ApiKey,
    AsyncSession,
    CashuTransaction,
    get_session,
    release_stale_reservations,
)
from .core.logging import get_logger
from .core.settings import settings
from .foreign_mint_swap import swap_enabled, swap_in_and_credit
from .lightning import lightning_router
from .wallet import (
    UntrustedSourceMintError,
    classify_redemption_error,
    credit_balance,
    is_trusted_source_mint,
    recieve_token,
    token_mint_url,
)

router = APIRouter()
balance_router = APIRouter(prefix="/v1/balance")

logger = get_logger(__name__)


async def get_key_from_header(
    authorization: Annotated[str, Header(...)],
    session: AsyncSession = Depends(get_session),
) -> ApiKey:
    if authorization.startswith("Bearer "):
        return await validate_bearer_key(authorization[7:], session)

    raise HTTPException(
        status_code=401,
        detail="Invalid authorization. Use 'Bearer <cashu-token>' or 'Bearer <api-key>'",
    )


async def get_balance_info(key: ApiKey, session: AsyncSession) -> dict:
    info = {
        "api_key": "sk-" + key.hashed_key,
        "balance": key.total_balance,
        "reserved": key.reserved_balance,
        "total_requests": key.total_requests,
        "total_spent": key.total_spent,
        "validity_date": key.validity_date,
    }

    return info


# TODO: remove this endpoint when frontend is updated
@router.get("/", include_in_schema=False)
async def account_info(
    key: ApiKey = Depends(get_key_from_header),
    session: AsyncSession = Depends(get_session),
) -> dict:
    return await get_balance_info(key, session)


# TODO: Implement POST /v1/wallet/create endpoint
# This endpoint should accept:
# - cashu_token (required): The eCash token to deposit
# - refund_lnurl (optional): LNURL for refunds (instead of refund_address in validate_bearer_key)
# - refund_expiry (optional): Expiry timestamp for the key (maps to key_expiry_time in validate_bearer_key)
# The endpoint should:
# 1. Create a new wallet/API key from the cashu_token
# 2. Store refund_lnurl and refund_expiry in the database
# 3. Return the API key (rstr_...) and balance
# Note: validate_bearer_key already supports refund_address and key_expiry_time params


class BalanceCreateRequest(BaseModel):
    initial_balance_token: str
    validity_date: int | None = None


async def _create_balance(
    initial_balance_token: str,
    validity_date: int | None,
    session: AsyncSession,
) -> dict:
    key = await validate_bearer_key(initial_balance_token, session)

    if validity_date is not None:
        key.validity_date = validity_date
        session.add(key)
        await session.commit()
        await session.refresh(key)

    return {
        "api_key": "sk-" + key.hashed_key,
        "balance": key.balance,
    }


@router.post("/create")
async def create_balance_from_body(
    payload: BalanceCreateRequest,
    session: AsyncSession = Depends(get_session),
) -> dict:
    return await _create_balance(
        payload.initial_balance_token,
        payload.validity_date,
        session,
    )


@router.get("/create")
async def create_balance(
    initial_balance_token: str,
    validity_date: int | None = None,
    session: AsyncSession = Depends(get_session),
) -> dict:
    return await _create_balance(
        initial_balance_token,
        validity_date,
        session,
    )


@router.get("/info")
async def wallet_info(
    key: ApiKey = Depends(get_key_from_header),
    session: AsyncSession = Depends(get_session),
) -> dict:
    return await get_balance_info(key, session)


class TopupRequest(BaseModel):
    cashu_token: str


def _error_chain(error: BaseException) -> list[dict[str, str]]:
    chain: list[dict[str, str]] = []
    current: BaseException | None = error
    seen: set[int] = set()
    while current is not None and id(current) not in seen:
        seen.add(id(current))
        chain.append({"type": type(current).__name__, "message": str(current)})
        current = current.__cause__ or current.__context__
    return chain


@router.post("/topup")
async def topup_wallet_endpoint(
    cashu_token: str | None = None,
    topup_request: TopupRequest | None = None,
    key: ApiKey = Depends(get_key_from_header),
    session: AsyncSession = Depends(get_session),
) -> dict[str, int]:
    billing_key = key

    if topup_request is not None:
        cashu_token = topup_request.cashu_token
    if cashu_token is None:
        raise HTTPException(status_code=400, detail="A cashu_token is required.")

    cashu_token = cashu_token.replace("\n", "").replace("\r", "").replace("\t", "")
    if len(cashu_token) < 10 or "cashu" not in cashu_token:
        raise HTTPException(status_code=400, detail="Invalid token format")

    source_mint = token_mint_url(cashu_token, "unknown")
    logger.info(
        "Cashu wallet top-up started",
        extra={
            "event": "cashu_topup_started",
            "source_mint": source_mint,
            "primary_mint": settings.primary_mint,
            "trusted_mints": settings.cashu_mints,
            "key_hash": billing_key.hashed_key[:8],
        },
    )
    try:
        if source_mint != "unknown" and not is_trusted_source_mint(source_mint):
            if not swap_enabled():
                raise UntrustedSourceMintError(f"Untrusted source mint: {source_mint}")
            # Top-up is the only entry point that swaps: the caller is already
            # waiting on a long operation here, unlike bearer auth or X-Cashu.
            amount_msats = await swap_in_and_credit(cashu_token, billing_key, session)
        else:
            amount_msats = await credit_balance(cashu_token, billing_key, session)
    except Exception as e:
        # Shared taxonomy so top-up matches the bearer/X-Cashu paths (503 for an
        # unreachable mint, 422 for fee/swap failures, 400 for token faults).
        classified = classify_redemption_error(e)
        if classified is None:
            logger.error(
                "Cashu wallet top-up failed with an unhandled error",
                extra={
                    "event": "cashu_topup_failed",
                    "source_mint": source_mint,
                    "primary_mint": settings.primary_mint,
                    "trusted_mints": settings.cashu_mints,
                    "error_chain": _error_chain(e),
                },
            )
            raise redemption_error_to_http_exception(e)
        error_type, status_code, message, error_code = classified
        logger.warning(
            "Cashu wallet top-up failed",
            extra={
                "event": "cashu_topup_failed",
                "source_mint": source_mint,
                "primary_mint": settings.primary_mint,
                "trusted_mints": settings.cashu_mints,
                "status_code": status_code,
                "error_type": error_type,
                "error_code": error_code,
                "error_chain": _error_chain(e),
            },
        )
        raise redemption_error_to_http_exception(e)

    logger.info(
        "Cashu wallet top-up completed",
        extra={
            "event": "cashu_topup_completed",
            "source_mint": source_mint,
            "credited_msats": amount_msats,
            "key_hash": billing_key.hashed_key[:8],
        },
    )
    return {"msats": amount_msats}


async def _lookup_key_no_create(
    bearer_value: str, session: AsyncSession
) -> ApiKey | None:
    """Look up an existing API key without creating one Used by the refund endpoint"""
    if bearer_value.startswith("sk-"):
        return await session.get(ApiKey, bearer_value[3:])
    if bearer_value.startswith("cashu"):
        hashed = hashlib.sha256(bearer_value.encode()).hexdigest()
        return await session.get(ApiKey, hashed)
    return None


async def _get_persisted_api_key_refund(
    key: ApiKey, session: AsyncSession, token: str | None = None
) -> dict[str, str] | None:
    query = select(CashuTransaction).where(
        CashuTransaction.api_key_hashed_key == key.hashed_key,
        CashuTransaction.type == "out",
        CashuTransaction.source == "apikey",
    )
    if token is not None:
        query = query.where(CashuTransaction.token == token)
    result = await session.exec(query.order_by(col(CashuTransaction.created_at).desc()))
    refund = result.first()
    if refund is None:
        return None
    if refund.swept:
        raise HTTPException(status_code=410, detail="Refund has been swept")

    refund.collected = True
    session.add(refund)
    await session.commit()

    persisted = {"token": refund.token}
    if refund.unit == "sat":
        persisted["sats"] = str(refund.amount)
    else:
        persisted["msats"] = str(refund.amount)
    return persisted


class RefundRequest(BaseModel):
    lightning_address: str | None = None


@router.post("/refund", response_model=None)
async def refund_wallet_endpoint(
    refund_request: RefundRequest | None = None,
    authorization: Annotated[str | None, Header()] = None,
    x_cashu: Annotated[str | None, Header()] = None,
    session: AsyncSession = Depends(get_session),
) -> JSONResponse | dict[str, str]:
    if x_cashu:
        # Find the "in" transaction by the original payment token
        in_tx_result = await session.exec(
            select(CashuTransaction).where(
                CashuTransaction.token == x_cashu,
                CashuTransaction.type == "in",
            )
        )
        in_tx = in_tx_result.first()
        if in_tx is None:
            raise HTTPException(status_code=404, detail="Refund not found")

        # Use the request_id to find the associated "out" (refund) transaction
        if in_tx.request_id is None:
            raise HTTPException(status_code=404, detail="Refund not found")

        out_tx_result = await session.exec(
            select(CashuTransaction).where(
                CashuTransaction.request_id == in_tx.request_id,
                CashuTransaction.type == "out",
            )
        )
        out_tx = out_tx_result.first()
        if out_tx is None:
            # The "in" row exists with a request_id, but the "out" (refund)
            # row hasn't been written yet — the upstream request is still in
            # flight and the refund will be minted once it completes. Tell the
            # client to retry instead of 404ing permanently (race condition
            # where /v1/wallet/refund is polled before the refund exists).
            logger.debug(
                "refund_wallet_endpoint: refund pending (in row exists, out row not yet created)",
                extra={"request_id": in_tx.request_id},
            )
            raise HTTPException(
                status_code=425,
                detail="Refund is pending; retry shortly.",
                headers={"Retry-After": "2"},
            )
        if out_tx.swept:
            raise HTTPException(status_code=410, detail="Refund has been swept")

        out_tx.collected = True
        session.add(out_tx)
        await session.commit()
        body: dict[str, str] = {"token": out_tx.token}
        if out_tx.unit == "sat":
            body["sats"] = str(out_tx.amount)
        else:
            body["msats"] = str(out_tx.amount)
        return JSONResponse(content=body, headers={"X-Cashu": out_tx.token})

    if authorization is None or not authorization.startswith("Bearer "):
        raise HTTPException(
            status_code=401,
            detail="Invalid authorization. Use 'Bearer <cashu-token>' or 'Bearer <api-key>'",
        )

    bearer_value: str = authorization[7:]
    key: ApiKey | None = await _lookup_key_no_create(bearer_value, session)
    if key is None:
        raise HTTPException(
            status_code=401,
            detail={
                "error": {
                    "message": "Key not found. Deposit first via /v1/wallet/create before requesting a refund.",
                    "type": "invalid_request_error",
                    "code": "key_not_found",
                }
            },
        )

    # Check for an open claim before any replay or destination lookup.
    if open_claim := await refund.latest_open(session, key):
        raise refund.refund_in_progress_error(open_claim)

    if key.total_balance <= 0:
        paid = await refund.latest_terminal(session, key)
        if paid and paid.method == "lightning":
            return refund.describe(paid)
        if paid and paid.token:
            # Match the ledger row to this claim's token, not the latest one.
            if persisted := await _get_persisted_api_key_refund(
                key, session, paid.token
            ):
                return persisted
            return refund.describe(paid)
        # Legacy payouts predate the claim row, so fall back to the ledger.
        if persisted := await _get_persisted_api_key_refund(key, session):
            return persisted
        if paid:
            return refund.describe(paid)
        if stuck := await refund.latest_stuck(session, key):
            raise refund.refund_in_progress_error(stuck)

    if key.reserved_balance > 0:
        # Release only durable reservations old enough to be stale. A newer
        # request on the same aggregate balance must remain reserved.
        await release_stale_reservations(
            session,
            settings.stale_reservation_timeout_seconds,
            key_hash=key.hashed_key,
        )
        await session.refresh(key)
        if key.reserved_balance > 0:
            raise HTTPException(
                status_code=400,
                detail="Cannot refund key. There are ongoing requests for this api key.",
            )
        logger.warning(
            "refund_wallet_endpoint: released stale reservation before refund",
            extra={
                "key_hash": key.hashed_key[:8],
                "stale_timeout_seconds": settings.stale_reservation_timeout_seconds,
            },
        )

    remaining_balance_msats: int = key.total_balance
    unit = refund.refund_unit(key)
    remaining_balance = refund.amount_in_unit(remaining_balance_msats, unit)

    if remaining_balance_msats > 0 and remaining_balance <= 0:
        raise HTTPException(status_code=400, detail="Balance too small to refund")
    elif remaining_balance <= 0:
        raise HTTPException(status_code=400, detail="No balance to refund")

    requested = refund_request.lightning_address if refund_request else None
    destination = requested or key.refund_address
    if destination:
        # Stored addresses can rot too; reject before any balance is debited.
        await refund.validate_lightning_destination(destination)

    claim = await refund.open_claim(
        session,
        key,
        method="lightning" if destination else "cashu",
        destination=destination,
    )
    return await refund.execute(session, claim)


@router.get("/history")
async def wallet_history(
    key: ApiKey = Depends(get_key_from_header),
    session: AsyncSession = Depends(get_session),
) -> dict[str, list[dict[str, str | int | bool | None]]]:
    result = await session.exec(
        select(CashuTransaction)
        .where(CashuTransaction.api_key_hashed_key == key.hashed_key)
        .order_by(col(CashuTransaction.created_at).desc())
    )
    transactions = result.all()
    return {
        "transactions": [
            {
                "id": tx.id,
                "type": tx.type,
                "source": tx.source,
                "amount": tx.amount,
                "unit": tx.unit,
                "mint_url": tx.mint_url,
                "created_at": tx.created_at,
                "collected": tx.collected,
                "swept": tx.swept,
            }
            for tx in transactions
        ]
    }


@router.post("/donate")
async def donate(token: str, ref: str | None = None) -> str:
    try:
        amount, unit, _ = await recieve_token(token)
        if ref:
            logger.info(
                "donation received", extra={"ref": ref, "amount": amount, "unit": unit}
            )
        return "Thanks!"
    except Exception:
        return "Invalid token."


@router.api_route(
    "/{path:path}",
    methods=["GET", "POST", "PUT", "DELETE"],
    include_in_schema=False,
    response_model=None,
)
async def wallet_catch_all(path: str) -> NoReturn:
    raise HTTPException(
        status_code=404, detail="Not found check /docs for available endpoints"
    )


balance_router.include_router(lightning_router, include_in_schema=False)
balance_router.include_router(router)

deprecated_wallet_router = APIRouter(prefix="/v1/wallet", include_in_schema=False)
deprecated_wallet_router.include_router(router)