routstr-core/routstr/core/db.py

import asyncio
import hashlib
import os
import pathlib
import sqlite3
import time
import uuid
from contextlib import asynccontextmanager
from enum import Enum
from typing import AsyncGenerator

from alembic import command
from alembic.config import Config
from alembic.util.exc import CommandError
from sqlalchemy import Index, UniqueConstraint, case, delete, event, or_
from sqlalchemy.engine import make_url
from sqlalchemy.exc import IntegrityError, OperationalError
from sqlalchemy.ext.asyncio import AsyncEngine
from sqlalchemy.ext.asyncio.engine import create_async_engine
from sqlalchemy.orm import aliased
from sqlmodel import Field, Relationship, SQLModel, col, func, select, update
from sqlmodel.ext.asyncio.session import AsyncSession

from .logging import get_logger
from .settings import settings

logger = get_logger(__name__)

DATABASE_URL = os.environ.get("DATABASE_URL", "sqlite+aiosqlite:///keys.db")


def create_db_engine(database_url: str = DATABASE_URL) -> AsyncEngine:
    """Build and instrument an async engine from environment-only settings."""
    url = make_url(database_url)
    backend = url.get_backend_name()
    is_sqlite = backend == "sqlite"
    is_memory_sqlite = is_sqlite and url.database in {None, "", ":memory:"}
    pool_pre_ping = settings.database_pool_pre_ping or not is_sqlite
    options: dict[str, int | float | bool] = {"pool_pre_ping": pool_pre_ping}
    if not is_memory_sqlite:
        options.update(
            pool_size=settings.database_pool_size,
            max_overflow=settings.database_max_overflow,
            pool_timeout=settings.database_pool_timeout,
            pool_recycle=settings.database_pool_recycle,
        )

    logger.info(
        "Database pool configured",
        extra={
            "database_url_backend": backend,
            "in_memory_sqlite": is_memory_sqlite,
            **options,
        },
    )
    created_engine = create_async_engine(database_url, echo=False, **options)
    hold_warn_seconds = settings.database_pool_hold_warn_seconds

    def record_pool_checkout(
        dbapi_connection: object, connection_record: object, proxy: object
    ) -> None:
        connection_record.info["routstr_checked_out_at"] = time.monotonic()  # type: ignore[attr-defined]

    def record_pool_checkin(
        dbapi_connection: object, connection_record: object
    ) -> None:
        checked_out_at = connection_record.info.pop(  # type: ignore[attr-defined]
            "routstr_checked_out_at", None
        )
        if checked_out_at is None:
            return
        held_seconds = time.monotonic() - checked_out_at
        if held_seconds >= hold_warn_seconds:
            logger.warning(
                "Database connection held longer than threshold",
                extra={
                    "held_seconds": round(held_seconds, 3),
                    "threshold_seconds": hold_warn_seconds,
                    "pool_status": created_engine.pool.status(),
                },
            )

    event.listen(created_engine.sync_engine, "checkout", record_pool_checkout)
    event.listen(created_engine.sync_engine, "checkin", record_pool_checkin)
    return created_engine


engine = create_db_engine()


class ApiKey(SQLModel, table=True):  # type: ignore
    __tablename__ = "api_keys"

    hashed_key: str = Field(primary_key=True)
    balance: int = Field(default=0, description="Balance in millisatoshis (msats)")
    reserved_balance: int = Field(
        default=0, description="Reserved balance in millisatoshis (msats)"
    )
    reserved_at: int | None = Field(
        default=None,
        description=(
            "Unix timestamp of the most recent balance reservation. Used to "
            "detect and release stale reservations (e.g. after client "
            "disconnects). NULL when no reservation has been made yet."
        ),
    )
    refund_address: str | None = Field(
        default=None,
        description="Lightning address to refund remaining balance after key expires",
    )
    key_expiry_time: int | None = Field(
        default=None,
        description="Unix-timestamp after which the cashu-token's balance gets refunded to the refund_address",
    )
    total_spent: int = Field(
        default=0, description="Total spent in millisatoshis (msats)"
    )
    total_requests: int = Field(default=0)
    created_at: int | None = Field(
        default_factory=lambda: int(time.time()),
        nullable=True,
        description=(
            "Unix timestamp when the key was created. Nullable: keys created "
            "before this column existed have no value and sort last."
        ),
    )
    refund_mint_url: str | None = Field(
        default=None,
        description="URL of the mint used to create the cashu-token",
    )
    refund_currency: str | None = Field(
        default=None,
        description="Currency of the cashu-token",
    )
    parent_key_hash: str | None = Field(
        default=None, foreign_key="api_keys.hashed_key", index=True
    )
    balance_limit: int | None = Field(
        default=None,
        description="Max spendable balance in msats for this key (mostly for child keys)",
    )
    balance_limit_reset: str | None = Field(
        default=None,
        description="Reset policy for balance limit (manual, daily, monthly, etc.)",
    )
    balance_limit_reset_date: int | None = Field(
        default=None,
        description="Unix timestamp of the last time the balance limit was reset",
    )
    validity_date: int | None = Field(
        default=None,
        description="Unix timestamp after which the key is no longer valid",
    )

    @property
    def total_balance(self) -> int:
        return self.balance - self.reserved_balance


async def reset_all_reserved_balances(session: AsyncSession) -> None:
    """Release every active durable reservation during explicit startup reset."""
    await session.exec(  # type: ignore[call-overload]
        update(ReservationRelease)
        .where(col(ReservationRelease.status) == "active")
        .values(status="released")
    )
    await session.exec(  # type: ignore[call-overload]
        update(ApiKey).values(reserved_balance=0, reserved_at=None)
    )
    await session.commit()
    logger.info("Reset reserved balances on startup")


async def release_stale_reservations(
    session: AsyncSession,
    max_age_seconds: int,
    *,
    key_hash: str | None = None,
) -> int:
    """Release stale durable reservations without touching newer reservations."""
    cutoff = int(time.time()) - max_age_seconds
    query = (
        select(ReservationRelease)
        .where(col(ReservationRelease.status) == "active")
        .where(col(ReservationRelease.created_at) < cutoff)
    )
    if key_hash is not None:
        query = query.where(
            or_(
                col(ReservationRelease.key_hash) == key_hash,
                col(ReservationRelease.billing_key_hash) == key_hash,
            )
        )
    reservations = (await session.exec(query)).all()
    released = 0

    for reservation in reservations:
        transition = await session.exec(  # type: ignore[call-overload]
            update(ReservationRelease)
            .where(col(ReservationRelease.id) == reservation.id)
            .where(col(ReservationRelease.status) == "active")
            .values(status="released")
        )
        if transition.rowcount != 1:
            continue

        values = {
            "reserved_balance": col(ApiKey.reserved_balance)
            - reservation.reserved_msats,
            "reserved_at": case(
                (
                    col(ApiKey.reserved_balance) - reservation.reserved_msats > 0,
                    col(ApiKey.reserved_at),
                ),
                else_=None,
            ),
        }
        parent_result = await session.exec(  # type: ignore[call-overload]
            update(ApiKey)
            .where(col(ApiKey.hashed_key) == reservation.billing_key_hash)
            .where(col(ApiKey.reserved_balance) >= reservation.reserved_msats)
            .values(**values)
        )
        if parent_result.rowcount != 1:
            await session.rollback()
            return 0

        if reservation.billing_key_hash != reservation.key_hash:
            child_result = await session.exec(  # type: ignore[call-overload]
                update(ApiKey)
                .where(col(ApiKey.hashed_key) == reservation.key_hash)
                .where(col(ApiKey.reserved_balance) >= reservation.reserved_msats)
                .values(**values)
            )
            if child_result.rowcount != 1:
                await session.rollback()
                return 0
        released += 1

    # Rolling upgrades can leave aggregate reservations created before durable
    # reservation rows existed. Release only stale aggregates that have no active
    # durable owner; targeted refund cleanup also heals legacy NULL timestamps.
    legacy_query = select(ApiKey).where(col(ApiKey.reserved_balance) > 0)
    if key_hash is None:
        legacy_query = legacy_query.where(col(ApiKey.reserved_at).is_not(None)).where(
            col(ApiKey.reserved_at) < cutoff
        )
    else:
        legacy_query = legacy_query.where(
            or_(
                col(ApiKey.hashed_key) == key_hash,
                col(ApiKey.parent_key_hash) == key_hash,
            )
        ).where(
            or_(col(ApiKey.reserved_at).is_(None), col(ApiKey.reserved_at) < cutoff)
        )

    for legacy_key in (await session.exec(legacy_query)).all():
        active_owner = (
            await session.exec(
                select(ReservationRelease.id)
                .where(col(ReservationRelease.status) == "active")
                .where(
                    or_(
                        col(ReservationRelease.key_hash) == legacy_key.hashed_key,
                        col(ReservationRelease.billing_key_hash)
                        == legacy_key.hashed_key,
                    )
                )
                .limit(1)
            )
        ).first()
        if active_owner is not None:
            continue
        legacy_key.reserved_balance = 0
        legacy_key.reserved_at = None
        session.add(legacy_key)
        released += 1

    await session.commit()
    if released:
        logger.warning(
            "Released stale reservations",
            extra={
                "released_reservations": released,
                "max_age_seconds": max_age_seconds,
            },
        )
    return released


async def prune_dead_api_keys(session: AsyncSession, min_age_seconds: int) -> int:
    """Delete dead parentless API keys; return the count removed.

    Dead = 0 balance/reservation/spend/requests, older than the grace period,
    no parent, no children, no pending invoice. Cashu rows are unlinked (not
    deleted) first to keep the audit trail.
    """
    cutoff = int(time.time()) - min_age_seconds

    child = aliased(ApiKey)
    has_children = (
        select(child.hashed_key).where(
            col(child.parent_key_hash) == col(ApiKey.hashed_key)
        )
    ).exists()
    pending_invoice = (
        select(LightningInvoice.id)
        .where(col(LightningInvoice.api_key_hash) == col(ApiKey.hashed_key))
        .where(col(LightningInvoice.status) == "pending")
    ).exists()

    eligible_hashes = (
        select(ApiKey.hashed_key)
        .where(col(ApiKey.balance) == 0)
        .where(col(ApiKey.reserved_balance) == 0)
        .where(col(ApiKey.total_spent) == 0)
        .where(col(ApiKey.total_requests) == 0)
        .where(col(ApiKey.parent_key_hash).is_(None))
        .where(
            (col(ApiKey.created_at).is_(None)) | (col(ApiKey.created_at) < cutoff)
        )
        .where(~pending_invoice)
        .where(~has_children)
    )

    # Unlink transactions rather than cascade-deleting them, so the financial
    # audit trail survives. The eligibility predicate is re-evaluated inside both
    # statements so a key that gained balance mid-run is left untouched.
    await session.exec(  # type: ignore[call-overload]
        update(CashuTransaction)
        .where(col(CashuTransaction.api_key_hashed_key).in_(eligible_hashes))
        .values(api_key_hashed_key=None)
    )

    result = await session.exec(  # type: ignore[call-overload]
        delete(ApiKey).where(col(ApiKey.hashed_key).in_(eligible_hashes))
    )
    await session.commit()

    pruned = int(result.rowcount or 0)
    logger.info(
        "Pruned dead API keys",
        extra={"pruned_keys": pruned, "min_age_seconds": min_age_seconds},
    )
    return pruned


class ModelRow(SQLModel, table=True):  # type: ignore
    __tablename__ = "models"
    id: str = Field(primary_key=True)
    upstream_provider_id: int = Field(
        primary_key=True, foreign_key="upstream_providers.id", ondelete="CASCADE"
    )
    name: str = Field()
    created: int = Field()
    description: str = Field()
    context_length: int = Field()
    architecture: str = Field()
    pricing: str = Field()
    sats_pricing: str | None = Field(default=None)
    per_request_limits: str | None = Field(default=None)
    top_provider: str | None = Field(default=None)
    canonical_slug: str | None = Field(default=None, description="Canonical model slug")
    alias_ids: str | None = Field(
        default=None, description="JSON array of model alias IDs"
    )
    enabled: bool = Field(default=True, description="Whether this model is enabled")
    forwarded_model_id: str | None = Field(
        default=None,
        description="Model ID to use when forwarding requests to upstream provider. Defaults to id if not set.",
    )
    upstream_provider: "UpstreamProviderRow" = Relationship(back_populates="models")


class LightningInvoice(SQLModel, table=True):  # type: ignore
    __tablename__ = "lightning_invoices"

    id: str = Field(primary_key=True, description="Unique invoice identifier")
    bolt11: str = Field(description="BOLT11 invoice string", unique=True)
    amount_sats: int = Field(description="Amount in satoshis")
    description: str = Field(description="Invoice description")
    payment_hash: str = Field(description="Payment hash for tracking", unique=True)
    status: str = Field(
        default="pending",
        description="pending, paid, expired, cancelled, reconciliation_required",
    )
    api_key_hash: str | None = Field(
        default=None, description="Associated API key hash for topup operations"
    )
    purpose: str = Field(description="create or topup")
    mint_url: str | None = Field(
        default=None, description="Mint URL where the quote was created (fallback tracking)"
    )
    created_at: int = Field(
        default_factory=lambda: int(time.time()), description="Unix timestamp"
    )
    expires_at: int = Field(description="Unix timestamp when invoice expires")
    paid_at: int | None = Field(default=None, description="Unix timestamp when paid")
    balance_limit: int | None = Field(
        default=None,
        description="Max spendable msats for the created key",
    )
    balance_limit_reset: str | None = Field(
        default=None,
        description="Reset policy for balance limit (daily, weekly, monthly)",
    )
    validity_date: int | None = Field(
        default=None,
        description="Unix timestamp after which the created key expires",
    )


class CashuTransaction(SQLModel, table=True):  # type: ignore
    __tablename__ = "cashu_transactions"

    id: str = Field(
        primary_key=True,
        default_factory=lambda: uuid.uuid4().hex,
        description="Unique transaction identifier",
    )
    token: str = Field(description="Serialized Cashu token")
    amount: int = Field(description="Amount in the token's unit")
    unit: str = Field(description="Token unit (sat or msat)")
    mint_url: str | None = Field(default=None, description="Mint URL for the token")
    type: str = Field(default="out", description="Transaction type: in or out")
    request_id: str | None = Field(default=None, description="Associated request ID")
    created_at: int = Field(
        default_factory=lambda: int(time.time()),
        description="Unix timestamp",
    )
    collected: bool = Field(default=False)
    swept: bool = Field(default=False)
    sweep_started_at: int | None = Field(
        default=None,
        description="Unix timestamp for a recoverable refund-sweep claim",
    )
    source: str = Field(
        default="x-cashu",
        description="Payment source: x-cashu or apikey",
    )
    api_key_hashed_key: str | None = Field(
        default=None,
        foreign_key="api_keys.hashed_key",
        index=True,
        description="Associated API key hash for wallet history",
    )


async def store_cashu_transaction(
    token: str,
    amount: int,
    unit: str,
    mint_url: str | None = None,
    typ: str = "out",
    request_id: str | None = None,
    collected: bool = False,
    created_at: int | None = None,
    source: str = "x-cashu",
    api_key_hashed_key: str | None = None,
    transaction_id: str | None = None,
    log_failure: bool = True,
) -> bool:
    try:
        async with create_session() as session:
            tx = CashuTransaction(
                id=transaction_id or uuid.uuid4().hex,
                token=token,
                amount=amount,
                unit=unit,
                mint_url=mint_url,
                type=typ,
                request_id=request_id,
                collected=collected,
                created_at=created_at or int(time.time()),
                source=source,
                api_key_hashed_key=api_key_hashed_key,
            )
            session.add(tx)
            await session.commit()
    except Exception:
        if log_failure:
            logger.critical(
                "Failed to store Cashu transaction",
                extra={"type": typ, "request_id": request_id, "source": source},
                exc_info=True,
            )
        raise
    return True


async def _cashu_transaction_exists(transaction_id: str) -> bool:
    async with create_session() as session:
        return await session.get(CashuTransaction, transaction_id) is not None


async def store_cashu_transaction_with_retry(
    token: str,
    amount: int,
    unit: str,
    mint_url: str | None = None,
    typ: str = "out",
    request_id: str | None = None,
    collected: bool = False,
    created_at: int | None = None,
    source: str = "x-cashu",
    api_key_hashed_key: str | None = None,
    max_attempts: int = 3,
) -> bool:
    """Retry a critical Cashu transaction write with bounded backoff."""
    transaction_id = hashlib.sha256(f"{typ}\0{token}".encode()).hexdigest()
    last_error: Exception | None = None
    for attempt in range(1, max_attempts + 1):
        try:
            return await store_cashu_transaction(
                token=token,
                amount=amount,
                unit=unit,
                mint_url=mint_url,
                typ=typ,
                request_id=request_id,
                collected=collected,
                created_at=created_at,
                source=source,
                api_key_hashed_key=api_key_hashed_key,
                transaction_id=transaction_id,
                log_failure=False,
            )
        except IntegrityError as error:
            try:
                if await _cashu_transaction_exists(transaction_id):
                    return True
            except Exception as lookup_error:
                last_error = lookup_error
            else:
                last_error = error
        except Exception as error:
            last_error = error

        if last_error is not None:
            if attempt == max_attempts:
                break
            delay = 0.25 * (2 ** (attempt - 1))
            logger.warning(
                "Cashu transaction storage failed; retrying",
                extra={
                    "type": typ,
                    "request_id": request_id,
                    "attempt": attempt,
                    "max_attempts": max_attempts,
                    "retry_delay_seconds": delay,
                },
            )
            await asyncio.sleep(delay)

    logger.critical(
        "Cashu transaction storage failed after bounded retries",
        extra={
            "type": typ,
            "request_id": request_id,
            "attempts": max_attempts,
            "error": str(last_error),
        },
    )
    if last_error is None:
        raise RuntimeError("Cashu transaction storage failed without an exception")
    raise last_error


class UpstreamProviderRow(SQLModel, table=True):  # type: ignore
    __tablename__ = "upstream_providers"
    __table_args__ = (
        UniqueConstraint(
            "base_url", "api_key", name="uq_upstream_providers_base_url_api_key"
        ),
    )
    id: int | None = Field(default=None, primary_key=True)
    slug: str | None = Field(
        default=None,
        unique=True,
        index=True,
        description="Stable external slug used for updates via API key.",
    )
    provider_type: str = Field(
        description="Provider type: custom, openai, anthropic, azure, openrouter, etc."
    )
    base_url: str = Field(description="Base URL of the upstream API")
    api_key: str = Field(description="API key for the upstream provider")
    api_version: str | None = Field(
        default=None, description="API version for Azure OpenAI"
    )
    enabled: bool = Field(default=True, description="Whether this provider is enabled")
    provider_fee: float = Field(
        default=1.01, description="Provider fee multiplier (default 1%)"
    )
    provider_settings: str | None = Field(
        default=None, description="JSON string for provider-specific settings"
    )
    models: list["ModelRow"] = Relationship(
        back_populates="upstream_provider",
        sa_relationship_kwargs={"cascade": "all, delete-orphan"},
    )


class ReservationRelease(SQLModel, table=True):  # type: ignore
    __tablename__ = "reservation_releases"
    __table_args__ = (
        Index("ix_reservation_releases_status_created_at", "status", "created_at"),
    )

    id: str = Field(primary_key=True)
    key_hash: str = Field(index=True)
    billing_key_hash: str = Field(index=True)
    reserved_msats: int
    status: str = Field(default="active")
    created_at: int = Field(default_factory=lambda: int(time.time()))


class RoutstrFee(SQLModel, table=True):  # type: ignore
    __tablename__ = "routstr_fees"
    id: int = Field(default=1, primary_key=True)
    accumulated_msats: int = Field(default=0)
    total_paid_msats: int = Field(default=0)
    last_paid_at: int | None = Field(default=None)
    payout_in_progress_msats: int = Field(default=0)
    payout_started_at: int | None = Field(default=None)


class NsecState(str, Enum):
    """Ownership state of the node's nsec — an explicit 3-state machine.

    The single ``encrypted_nsec`` column cannot distinguish "never migrated" from
    "intentionally cleared" (both leave it empty), which let a cleared identity be
    resurrected from a stale legacy ``NSEC``. This names the three states so the
    bootstrap branches on ownership rather than inferring it:

    * ``legacy`` — the vault has not taken ownership; a plaintext ``NSEC`` (env or
      old settings blob) may still exist and should be migrated in once.
    * ``encrypted`` — the vault owns a ciphertext; decrypt it, never re-read env.
    * ``cleared`` — the vault owns it but the operator emptied it; stay empty,
      never re-import from a stale legacy copy.
    """

    legacy = "legacy"
    encrypted = "encrypted"
    cleared = "cleared"


class Secret(SQLModel, table=True):  # type: ignore
    """Node-level secrets, stored encrypted/hashed at rest (singleton, id=1).

    The asymmetric column names document the encoding: ``_hash`` is one-way
    (scrypt, verify only) while ``encrypted_`` is reversible (Fernet). Per-provider
    upstream keys live on ``upstream_providers``, not here. See ``routstr.core.vault``.
    """

    __tablename__ = "secrets"
    id: int = Field(default=1, primary_key=True)
    admin_password_hash: str | None = Field(default=None)
    encrypted_nsec: str | None = Field(default=None)
    nsec_state: NsecState = Field(default=NsecState.legacy)
    updated_at: int | None = Field(default=None)


class CliToken(SQLModel, table=True):  # type: ignore
    """Long-lived authorization token for CLI/agent use against admin endpoints."""

    __tablename__ = "cli_tokens"
    id: str = Field(
        primary_key=True, default_factory=lambda: uuid.uuid4().hex
    )
    token: str = Field(unique=True, index=True, description="Bearer token value")
    name: str = Field(description="Human-readable label for this token")
    created_at: int = Field(default_factory=lambda: int(time.time()))
    last_used_at: int | None = Field(default=None)
    expires_at: int | None = Field(
        default=None, description="Optional expiry unix timestamp; null = never expires"
    )


async def accumulate_routstr_fee(session: AsyncSession, amount_msats: int) -> None:
    stmt = (
        update(RoutstrFee)
        .where(col(RoutstrFee.id) == 1)
        .values(accumulated_msats=RoutstrFee.accumulated_msats + amount_msats)
    )
    result = await session.exec(stmt)  # type: ignore[call-overload]
    if result.rowcount == 0:
        session.add(RoutstrFee(id=1, accumulated_msats=amount_msats))
    await session.commit()


async def get_routstr_fee(session: AsyncSession) -> RoutstrFee:
    fee = await session.get(RoutstrFee, 1)
    if fee is None:
        fee = RoutstrFee(id=1, accumulated_msats=0, total_paid_msats=0)
        session.add(fee)
        await session.commit()
        await session.refresh(fee)
    return fee


async def get_secret(session: AsyncSession) -> Secret:
    secret = await session.get(Secret, 1)
    if secret is None:
        secret = Secret(id=1)
        session.add(secret)
        try:
            await session.commit()
        except IntegrityError:
            # Another worker created the singleton row between our read and
            # insert (multiple workers booting against one shared DB). Roll back
            # and read the row they committed instead of failing startup.
            await session.rollback()
            secret = await session.get(Secret, 1)
            if secret is None:
                raise
            return secret
        await session.refresh(secret)
    return secret


async def set_admin_password(session: AsyncSession, password: str) -> None:
    """Store the admin password as a one-way hash on the Secret singleton."""
    from .vault import hash_password

    secret = await get_secret(session)
    secret.admin_password_hash = hash_password(password)
    secret.updated_at = int(time.time())
    session.add(secret)
    await session.commit()


async def set_nsec(session: AsyncSession, nsec: str) -> None:
    """Store the node's nsec, Fernet-encrypted, on the Secret singleton.

    An empty string clears it (the node then holds no Nostr identity and signs
    no events). Either way the vault now owns the nsec, so the state moves off
    ``legacy``: a cleared identity (``cleared``) must not be resurrected from a
    stale legacy ``NSEC`` on the next boot.
    """
    from .vault import encrypt

    secret = await get_secret(session)
    secret.encrypted_nsec = encrypt(nsec) if nsec else None
    secret.nsec_state = NsecState.encrypted if nsec else NsecState.cleared
    secret.updated_at = int(time.time())
    session.add(secret)
    await session.commit()


async def reset_routstr_fee(session: AsyncSession, paid_msats: int) -> bool:
    """Checkpoint a fee payout before making the external payment."""
    stmt = (
        update(RoutstrFee)
        .where(col(RoutstrFee.id) == 1)
        .where(col(RoutstrFee.payout_in_progress_msats) == 0)
        .where(col(RoutstrFee.accumulated_msats) >= paid_msats)
        .values(
            accumulated_msats=RoutstrFee.accumulated_msats - paid_msats,
            payout_in_progress_msats=paid_msats,
            payout_started_at=int(time.time()),
        )
    )
    result = await session.exec(stmt)  # type: ignore[call-overload]
    await session.commit()
    return result.rowcount == 1


async def complete_routstr_fee_payout(
    session: AsyncSession, paid_msats: int
) -> bool:
    """Mark a checkpointed payout complete after the external payment succeeds."""
    stmt = (
        update(RoutstrFee)
        .where(col(RoutstrFee.id) == 1)
        .where(col(RoutstrFee.payout_in_progress_msats) == paid_msats)
        .values(
            payout_in_progress_msats=0,
            payout_started_at=None,
            total_paid_msats=RoutstrFee.total_paid_msats + paid_msats,
            last_paid_at=int(time.time()),
        )
    )
    result = await session.exec(stmt)  # type: ignore[call-overload]
    await session.commit()
    return result.rowcount == 1


async def total_user_liability(db_session: AsyncSession) -> int:
    """Return all outstanding API-key balances in millisatoshis."""
    result = await db_session.exec(select(func.sum(ApiKey.balance)))
    return int(result.one() or 0)


async def balance_for_mint_and_unit(
    db_session: AsyncSession, mint_url: str, unit: str
) -> int:
    """Return the user liability for one mint and unit in millisatoshis."""
    result = await db_session.exec(
        select(func.sum(ApiKey.balance)).where(
            col(ApiKey.refund_mint_url) == mint_url,
            col(ApiKey.refund_currency) == unit,
        )
    )
    return int(result.one() or 0)


async def balances_by_mint_and_unit(
    db_session: AsyncSession, mint_urls: list[str], units: list[str]
) -> dict[tuple[str, str], int]:
    """Return requested user liabilities grouped by mint and unit."""
    if not mint_urls or not units:
        return {}
    query = (
        select(
            col(ApiKey.refund_mint_url),
            col(ApiKey.refund_currency),
            func.sum(ApiKey.balance),
        )
        .where(
            col(ApiKey.refund_mint_url).in_(mint_urls),
            col(ApiKey.refund_currency).in_(units),
        )
        .group_by(col(ApiKey.refund_mint_url), col(ApiKey.refund_currency))
    )
    result = await db_session.exec(query)
    return {
        (mint_url, unit): int(balance or 0)
        for mint_url, unit, balance in result.all()
        if mint_url is not None and unit is not None
    }


async def init_db() -> None:
    """Initializes the database and creates tables if they don't exist."""
    async with engine.begin() as conn:
        if DATABASE_URL.startswith("sqlite"):
            await conn.exec_driver_sql("PRAGMA journal_mode=WAL")
        await conn.run_sync(SQLModel.metadata.create_all)


async def get_session() -> AsyncGenerator[AsyncSession, None]:
    async with AsyncSession(engine, expire_on_commit=False) as session:
        yield session


@asynccontextmanager
async def create_session() -> AsyncGenerator[AsyncSession, None]:
    async with AsyncSession(engine, expire_on_commit=False) as session:
        yield session


def fix_cashu_migrations() -> None:
    """
    Fixes Cashu wallet migrations that are not idempotent.
    This specifically addresses the 'duplicate column name: public_keys' error
    in the keysets table of Cashu's internal SQLite databases.
    """
    project_root = pathlib.Path(__file__).resolve().parents[2]
    wallet_dir = project_root / ".wallet"

    if not wallet_dir.exists() or not wallet_dir.is_dir():
        return

    logger.info("Checking Cashu wallet databases for migration idempotency")

    for db_file in wallet_dir.glob("*.sqlite3"):
        try:
            conn = sqlite3.connect(db_file)
            cursor = conn.cursor()

            # Check if keysets table exists
            cursor.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='keysets'"
            )
            if not cursor.fetchone():
                conn.close()
                continue

            # Check if public_keys column exists
            cursor.execute("PRAGMA table_info(keysets)")
            columns = [info[1] for info in cursor.fetchall()]

            if "public_keys" not in columns:
                logger.info(f"Adding missing public_keys column to {db_file.name}")
                cursor.execute("ALTER TABLE keysets ADD COLUMN public_keys TEXT")
                conn.commit()

            conn.close()
        except Exception as e:
            logger.warning(f"Could not check/fix Cashu database {db_file}: {e}")


def _clear_alembic_version() -> None:
    """Clear the alembic_version table so stamp/upgrade can proceed."""
    sync_url = DATABASE_URL.replace("+aiosqlite", "")
    from sqlalchemy import create_engine, text

    eng = create_engine(sync_url)
    with eng.begin() as conn:
        conn.execute(text("DELETE FROM alembic_version"))
    eng.dispose()


def run_migrations() -> None:
    """Run Alembic migrations programmatically."""
    try:
        # Run Cashu migration fix first
        fix_cashu_migrations()

        # Get the path to the alembic.ini file
        project_root = pathlib.Path(__file__).resolve().parents[2]
        alembic_ini_path = project_root / "alembic.ini"

        if not alembic_ini_path.exists():
            raise FileNotFoundError(
                f"Alembic configuration file not found at {alembic_ini_path}"
            )

        # Create Alembic config object
        alembic_cfg = Config(str(alembic_ini_path))

        # Set the database URL in the config
        alembic_cfg.set_main_option("sqlalchemy.url", DATABASE_URL)

        try:
            command.upgrade(alembic_cfg, "head")
        except CommandError as e:
            if "Can't locate revision" in str(e):
                logger.warning(
                    "Database stamped with unknown revision (likely from another branch). "
                    "Re-stamping to current head.",
                    extra={"error": str(e)},
                )
                _clear_alembic_version()
                command.stamp(alembic_cfg, "head")
            else:
                raise
        except OperationalError as e:
            if "duplicate column name" in str(e).lower():
                logger.warning(
                    "Migration hit a column that already exists (likely added via "
                    "create_all on another branch). Stamping to current head.",
                    extra={"error": str(e)},
                )
                _clear_alembic_version()
                command.stamp(alembic_cfg, "head")
            else:
                raise

        logger.info("Database migrations completed successfully")

    except Exception as e:
        logger.error(
            "Database migration failed",
            extra={"error": str(e), "error_type": type(e).__name__},
        )
        raise