routstr-core/routstr/core/db.py

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

from alembic import command
from alembic.config import Config
from alembic.util.exc import CommandError
from sqlalchemy import UniqueConstraint, delete
from sqlalchemy.exc import OperationalError
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

logger = get_logger(__name__)

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


engine = create_async_engine(DATABASE_URL, echo=False)  # echo=True for debugging SQL


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:
    stmt = update(ApiKey).values(reserved_balance=0, reserved_at=None)
    await session.exec(stmt)  # type: ignore[call-overload]
    await session.commit()
    logger.info("Reset reserved balances on startup")


async def release_stale_reservations(
    session: AsyncSession, max_age_seconds: int
) -> int:
    """Release reservations whose last reserve is older than max_age_seconds.
    """
    cutoff = int(time.time()) - max_age_seconds
    stmt = (
        update(ApiKey)
        .where(col(ApiKey.reserved_balance) > 0)
        .where(col(ApiKey.reserved_at).is_not(None))
        .where(col(ApiKey.reserved_at) < cutoff)
        .values(reserved_balance=0, reserved_at=None)
    )
    result = await session.exec(stmt)  # type: ignore[call-overload]
    await session.commit()
    released = int(result.rowcount or 0)
    if released:
        logger.warning(
            "Released stale balance reservations",
            extra={"released_keys": 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"
    )
    api_key_hash: str | None = Field(
        default=None, description="Associated API key hash for topup operations"
    )
    purpose: str = Field(description="create or topup")
    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)
    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,
) -> bool:
    try:
        async with create_session() as session:
            tx = CashuTransaction(
                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:
        logger.critical(
            "Failed to store Cashu transaction",
            extra={"type": typ, "request_id": request_id, "source": source},
            exc_info=True,
        )
        raise
    return True


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."""
    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,
            )
        except Exception as error:
            last_error = error
            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 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)


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 reset_routstr_fee(session: AsyncSession, paid_msats: int) -> None:
    stmt = (
        update(RoutstrFee)
        .where(col(RoutstrFee.id) == 1)
        .values(
            accumulated_msats=RoutstrFee.accumulated_msats - paid_msats,
            total_paid_msats=RoutstrFee.total_paid_msats + paid_msats,
            last_paid_at=int(time.time()),
        )
    )
    await session.exec(stmt)  # type: ignore[call-overload]
    await session.commit()


async def balances_for_mint_and_unit(
    db_session: AsyncSession, mint_url: str, unit: str
) -> int:
    query = select(func.sum(ApiKey.balance)).where(
        ApiKey.refund_mint_url == mint_url, ApiKey.refund_currency == unit
    )
    result = await db_session.exec(query)
    return result.one() or 0


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