routstr-core/routstr/lightning.py

import asyncio
import hashlib
import secrets
import time

from fastapi import APIRouter, Depends, Header, HTTPException
from pydantic import BaseModel, Field
from sqlmodel import col, select, update
from sqlmodel.ext.asyncio.session import AsyncSession

from .core.db import ApiKey, LightningInvoice, create_session, get_session
from .core.logging import get_logger
from .core.settings import settings
from .wallet import get_wallet

logger = get_logger(__name__)

lightning_router = APIRouter(prefix="/lightning")


class InvoiceCreateRequest(BaseModel):
    amount_sats: int = Field(gt=0, le=1_000_000, description="Amount in satoshis")
    purpose: str = Field(
        default="create",
        description="create or topup",
        pattern="^(create|topup)$",
    )
    api_key: str | None = Field(
        default=None,
        description="Deprecated: legacy field for topup. Prefer Authorization header.",
    )
    balance_limit: int | None = Field(default=None)
    balance_limit_reset: str | None = Field(default=None)
    validity_date: int | None = Field(default=None)


def _extract_bearer_api_key(authorization: str | None) -> str | None:
    if not authorization:
        return None
    token = authorization.strip()
    if token.lower().startswith("bearer "):
        token = token[7:].strip()
    return token or None


class InvoiceCreateResponse(BaseModel):
    invoice_id: str
    bolt11: str
    amount_sats: int
    expires_at: int
    payment_hash: str


class InvoiceStatusResponse(BaseModel):
    status: str
    api_key: str | None = None
    amount_sats: int
    paid_at: int | None = None
    created_at: int
    expires_at: int


class InvoiceRecoverRequest(BaseModel):
    bolt11: str = Field(description="BOLT11 invoice string")


async def generate_lightning_invoice(
    amount_sats: int, description: str
) -> tuple[str, str]:
    wallet = await get_wallet(settings.primary_mint, "sat")
    quote = await wallet.request_mint(amount_sats)
    return quote.request, quote.quote


def generate_invoice_id() -> str:
    return secrets.token_urlsafe(16)


@lightning_router.post("/invoice", response_model=InvoiceCreateResponse)
async def create_invoice(
    request: InvoiceCreateRequest,
    authorization: str | None = Header(default=None),
    session: AsyncSession = Depends(get_session),
) -> InvoiceCreateResponse:
    api_key_token = _extract_bearer_api_key(authorization) or request.api_key

    if request.purpose == "topup":
        if not api_key_token:
            raise HTTPException(
                status_code=401,
                detail="Authorization bearer api key is required for topup",
            )
        if not api_key_token.startswith("sk-"):
            raise HTTPException(status_code=400, detail="Invalid API key format")

        api_key = await session.get(ApiKey, api_key_token[3:])
        if not api_key:
            raise HTTPException(status_code=404, detail="API key not found")

    try:
        description = f"Routstr {request.purpose} {request.amount_sats} sats"
        bolt11, payment_hash = await generate_lightning_invoice(
            request.amount_sats, description
        )

        invoice_id = generate_invoice_id()
        expires_at = int(time.time()) + 3600  # 1 hour expiry

        invoice = LightningInvoice(
            id=invoice_id,
            bolt11=bolt11,
            amount_sats=request.amount_sats,
            description=description,
            payment_hash=payment_hash,
            status="pending",
            api_key_hash=api_key_token[3:] if api_key_token else None,
            purpose=request.purpose,
            balance_limit=request.balance_limit,
            balance_limit_reset=request.balance_limit_reset,
            validity_date=request.validity_date,
            expires_at=expires_at,
        )

        session.add(invoice)
        await session.commit()

        logger.info(
            "Lightning invoice created",
            extra={
                "invoice_id": invoice_id,
                "amount_sats": request.amount_sats,
                "purpose": request.purpose,
                "expires_at": expires_at,
            },
        )

        return InvoiceCreateResponse(
            invoice_id=invoice_id,
            bolt11=bolt11,
            amount_sats=request.amount_sats,
            expires_at=expires_at,
            payment_hash=payment_hash,
        )

    except Exception as e:
        logger.error(f"Failed to create Lightning invoice: {e}")
        raise HTTPException(
            status_code=500, detail="Failed to create Lightning invoice"
        )


@lightning_router.get(
    "/invoice/{invoice_id}/status", response_model=InvoiceStatusResponse
)
async def get_invoice_status(
    invoice_id: str,
    session: AsyncSession = Depends(get_session),
) -> InvoiceStatusResponse:
    invoice = await session.get(LightningInvoice, invoice_id)
    if not invoice:
        raise HTTPException(status_code=404, detail="Invoice not found")

    if invoice.status == "pending":
        await check_invoice_payment(invoice, session)

    if invoice.status == "pending" and int(time.time()) > invoice.expires_at:
        invoice.status = "expired"
        await session.commit()

    api_key = None
    if invoice.status == "paid" and invoice.purpose == "create":
        if invoice.api_key_hash:
            api_key = f"sk-{invoice.api_key_hash}"
    elif (
        invoice.status == "paid" and invoice.purpose == "topup" and invoice.api_key_hash
    ):
        api_key = f"sk-{invoice.api_key_hash}"

    return InvoiceStatusResponse(
        status=invoice.status,
        api_key=api_key,
        amount_sats=invoice.amount_sats,
        paid_at=invoice.paid_at,
        created_at=invoice.created_at,
        expires_at=invoice.expires_at,
    )


@lightning_router.post("/recover", response_model=InvoiceStatusResponse)
async def recover_invoice(
    request: InvoiceRecoverRequest,
    session: AsyncSession = Depends(get_session),
) -> InvoiceStatusResponse:
    result = await session.exec(
        select(LightningInvoice).where(LightningInvoice.bolt11 == request.bolt11)
    )
    invoice = result.first()

    if not invoice:
        raise HTTPException(status_code=404, detail="Invoice not found")

    if invoice.status == "pending":
        await check_invoice_payment(invoice, session)

    api_key = None
    if invoice.status == "paid":
        if invoice.purpose == "create" and invoice.api_key_hash:
            api_key = f"sk-{invoice.api_key_hash}"
        elif invoice.purpose == "topup" and invoice.api_key_hash:
            api_key = f"sk-{invoice.api_key_hash}"

    return InvoiceStatusResponse(
        status=invoice.status,
        api_key=api_key,
        amount_sats=invoice.amount_sats,
        paid_at=invoice.paid_at,
        created_at=invoice.created_at,
        expires_at=invoice.expires_at,
    )


async def check_invoice_payment(
    invoice: LightningInvoice, session: AsyncSession
) -> None:
    minted = False
    invoice_id = invoice.id
    invoice_purpose = invoice.purpose
    try:
        # A preceding invoice lookup starts a transaction. End it before the
        # potentially slow mint request so it cannot pin a pool connection.
        await session.commit()

        wallet = await get_wallet(settings.primary_mint, "sat")
        mint_status = await wallet.get_mint_quote(invoice.payment_hash)
        if not mint_status.paid:
            return

        # The mint enforces single-use quotes, so a concurrent checker that
        # races us here fails inside wallet.mint rather than double-minting.
        await wallet.mint(invoice.amount_sats, quote_id=invoice.payment_hash)
        minted = True

        if invoice_purpose == "create":
            api_key = await _create_api_key_record(invoice, session)
            invoice.api_key_hash = api_key.hashed_key
        elif invoice_purpose == "topup":
            await _credit_topup_record(invoice, session)

        # Conditional transition guards against double-credit: the credit
        # above and this status flip commit atomically, and a lost race
        # rolls both back.
        paid_at = int(time.time())
        finalized = await session.exec(  # type: ignore[call-overload]
            update(LightningInvoice)
            .where(
                col(LightningInvoice.id) == invoice_id,
                col(LightningInvoice.status) == "pending",
            )
            .values(
                status="paid",
                paid_at=paid_at,
                api_key_hash=invoice.api_key_hash,
            )
        )
        if finalized.rowcount != 1:
            await session.rollback()
            await session.refresh(invoice)
            return
        await session.commit()
        invoice.status = "paid"
        invoice.paid_at = paid_at

        logger.info(
            "Lightning invoice paid",
            extra={
                "invoice_id": invoice_id,
                "amount_sats": invoice.amount_sats,
                "purpose": invoice_purpose,
                "api_key_hash": invoice.api_key_hash[:8] + "..."
                if invoice.api_key_hash
                else None,
            },
        )
    except BaseException as e:
        # BaseException so task cancellation (e.g. client disconnect) after a
        # successful mint still triggers the reconciliation alert.
        try:
            await asyncio.shield(session.rollback())
        except Exception:
            logger.warning("Rollback failed during invoice check cleanup")
        if minted:
            logger.critical(
                "Invoice mint succeeded but DB finalization failed; reconciliation required",
                extra={"invoice_id": invoice_id, "purpose": invoice_purpose},
            )
        if not isinstance(e, Exception):
            raise
        logger.error(f"Failed to check invoice payment: {e}")


async def _create_api_key_record(
    invoice: LightningInvoice, session: AsyncSession
) -> ApiKey:
    dummy_token = f"invoice-{invoice.id}-{invoice.payment_hash}"
    hashed_key = hashlib.sha256(dummy_token.encode()).hexdigest()
    api_key = ApiKey(
        hashed_key=hashed_key,
        balance=invoice.amount_sats * 1000,
        refund_currency="sat",
        refund_mint_url=settings.primary_mint,
        balance_limit=invoice.balance_limit,
        balance_limit_reset=invoice.balance_limit_reset,
        validity_date=invoice.validity_date,
    )
    session.add(api_key)
    await session.flush()
    return api_key


async def _credit_topup_record(
    invoice: LightningInvoice, session: AsyncSession
) -> None:
    if not invoice.api_key_hash:
        raise ValueError("No API key associated with topup invoice")
    credited = await session.exec(  # type: ignore[call-overload]
        update(ApiKey)
        .where(col(ApiKey.hashed_key) == invoice.api_key_hash)
        .values(balance=col(ApiKey.balance) + invoice.amount_sats * 1000)
    )
    if credited.rowcount != 1:
        raise ValueError("Associated API key not found")


async def create_api_key_from_invoice(
    invoice: LightningInvoice, session: AsyncSession
) -> ApiKey:
    wallet = await get_wallet(settings.primary_mint, "sat")
    await wallet.mint(invoice.amount_sats, quote_id=invoice.payment_hash)
    return await _create_api_key_record(invoice, session)


async def topup_api_key_from_invoice(
    invoice: LightningInvoice, session: AsyncSession
) -> None:
    wallet = await get_wallet(settings.primary_mint, "sat")
    await wallet.mint(invoice.amount_sats, quote_id=invoice.payment_hash)
    await _credit_topup_record(invoice, session)


INVOICE_WATCH_INTERVAL_SECONDS = 5
INVOICE_WATCH_BATCH_LIMIT = 100


async def periodic_invoice_watcher() -> None:
    """Background task: detect paid Lightning invoices and credit balances.

    Removes the need for clients to poll the status endpoint after paying.
    """
    while True:
        try:
            async with create_session() as session:
                now = int(time.time())
                result = await session.exec(
                    select(LightningInvoice)
                    .where(
                        LightningInvoice.status == "pending",
                        col(LightningInvoice.expires_at) > now,
                    )
                    .limit(INVOICE_WATCH_BATCH_LIMIT)
                )
                pending = result.all()
                for invoice in pending:
                    try:
                        await check_invoice_payment(invoice, session)
                    except Exception as e:
                        logger.error(
                            "Invoice watcher failed for invoice",
                            extra={"invoice_id": invoice.id, "error": str(e)},
                        )
        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.error(f"Invoice watcher loop error: {e}")

        await asyncio.sleep(INVOICE_WATCH_INTERVAL_SECONDS)