exec-mcp/src/exec_mcp/repo.py

from __future__ import annotations

import json
import secrets
from datetime import datetime, timezone
from sqlite3 import IntegrityError
from uuid import uuid4

from fastapi import HTTPException

from .storage import get_connection


def utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def list_accounts(*, venue: str | None = None, enabled_only: bool = True) -> list[dict]:
    query = "SELECT id, display_name, venue, venue_account_ref, description, enabled, metadata_json, created_at, updated_at FROM accounts"
    clauses: list[str] = []
    params: list[str] = []

    if venue:
        clauses.append("venue = ?")
        params.append(venue)
    if enabled_only:
        clauses.append("enabled = 1")

    if clauses:
        query += " WHERE " + " AND ".join(clauses)
    query += " ORDER BY created_at ASC"

    with get_connection() as conn:
        rows = conn.execute(query, tuple(params)).fetchall()

    return [
        {
            "id": row["id"],
            "display_name": row["display_name"],
            "venue": row["venue"],
            "venue_account_ref": row["venue_account_ref"],
            "description": row["description"],
            "enabled": bool(row["enabled"]),
            "metadata": row["metadata_json"],
            "created_at": row["created_at"],
            "updated_at": row["updated_at"],
        }
        for row in rows
    ]


def get_account(account_id: str) -> dict:
    with get_connection() as conn:
        row = conn.execute(
            "SELECT id, display_name, venue, venue_account_ref, description, enabled, metadata_json, created_at, updated_at FROM accounts WHERE id = ?",
            (account_id,),
        ).fetchone()
    if row is None:
        raise HTTPException(status_code=404, detail="account not found")
    return {
        "id": row["id"],
        "display_name": row["display_name"],
        "venue": row["venue"],
        "venue_account_ref": row["venue_account_ref"],
        "description": row["description"],
        "enabled": bool(row["enabled"]),
        "metadata": row["metadata_json"],
        "created_at": row["created_at"],
        "updated_at": row["updated_at"],
    }


def get_account_secrets(account_id: str) -> dict:
    with get_connection() as conn:
        row = conn.execute(
            "SELECT api_key, api_secret FROM account_secrets WHERE account_id = ?",
            (account_id,),
        ).fetchone()
    if row is None:
        raise HTTPException(status_code=404, detail="account secrets not found")
    return {"api_key": row["api_key"], "api_secret": row["api_secret"]}


def _generate_account_id(length: int = 12) -> str:
    alphabet = "abcdefghijklmnopqrstuvwxyz0123456789"
    while True:
        candidate = "".join(secrets.choice(alphabet) for _ in range(length))
        with get_connection() as conn:
            exists = conn.execute("SELECT 1 FROM accounts WHERE id = ?", (candidate,)).fetchone()
        if exists is None:
            return candidate


def create_account(*, display_name: str, venue: str, venue_account_ref: str, api_key: str, api_secret: str, description: str | None = None, enabled: bool = True) -> dict:
    account_id = _generate_account_id()
    now = utc_now_iso()
    with get_connection() as conn:
        try:
            conn.execute(
                """
                INSERT INTO accounts (id, display_name, venue, venue_account_ref, description, enabled, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (account_id, display_name, venue, venue_account_ref, description, int(enabled), now, now),
            )
            conn.execute(
                """
                INSERT INTO account_secrets (account_id, api_key, api_secret, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?)
                """,
                (account_id, api_key, api_secret, now, now),
            )
            conn.commit()
        except IntegrityError as exc:
            conn.rollback()
            if "api_key" in str(exc).lower():
                raise HTTPException(status_code=409, detail="api key already exists") from exc
            raise
    return {"id": account_id, "display_name": display_name, "venue": venue}


def update_account(*, account_id: str, display_name: str | None = None, description: str | None = None, enabled: bool | None = None) -> dict:
    now = utc_now_iso()
    with get_connection() as conn:
        row = conn.execute("SELECT id FROM accounts WHERE id = ?", (account_id,)).fetchone()
        if row is None:
            raise HTTPException(status_code=404, detail="account not found")
        if display_name is not None:
            conn.execute("UPDATE accounts SET display_name = ?, updated_at = ? WHERE id = ?", (display_name, now, account_id))
        if description is not None:
            conn.execute("UPDATE accounts SET description = ?, updated_at = ? WHERE id = ?", (description, now, account_id))
        if enabled is not None:
            conn.execute("UPDATE accounts SET enabled = ?, updated_at = ? WHERE id = ?", (int(enabled), now, account_id))
        conn.commit()
    return {"id": account_id, "updated": True}


def delete_account(*, account_id: str) -> dict:
    with get_connection() as conn:
        deleted = conn.execute("DELETE FROM accounts WHERE id = ?", (account_id,)).rowcount
        conn.commit()
    if not deleted:
        raise HTTPException(status_code=404, detail="account not found")
    return {"id": account_id, "deleted": True}


def cache_get(cache_key: str) -> dict | None:
    now = utc_now_iso()
    with get_connection() as conn:
        row = conn.execute(
            "SELECT payload_json FROM api_cache WHERE cache_key = ? AND expires_at > ?",
            (cache_key, now),
        ).fetchone()
    if row is None:
        return None
    return json.loads(row["payload_json"])


def cache_put(cache_key: str, payload: dict, ttl_seconds: int) -> None:
    now_dt = datetime.now(timezone.utc)
    fetched_at = now_dt.isoformat()
    expires_at = datetime.fromtimestamp(now_dt.timestamp() + ttl_seconds, tz=timezone.utc).isoformat()
    with get_connection() as conn:
        conn.execute(
            """
            INSERT INTO api_cache (cache_key, payload_json, fetched_at, expires_at)
            VALUES (?, ?, ?, ?)
            ON CONFLICT(cache_key) DO UPDATE SET
                payload_json=excluded.payload_json,
                fetched_at=excluded.fetched_at,
                expires_at=excluded.expires_at
            """,
            (cache_key, json.dumps(payload), fetched_at, expires_at),
        )
        conn.commit()


def cache_delete(cache_key: str) -> None:
    with get_connection() as conn:
        conn.execute("DELETE FROM api_cache WHERE cache_key = ?", (cache_key,))
        conn.commit()


def has_account_balance_snapshot(*, account_id: str, source_kind: str, source_ref: str) -> bool:
    with get_connection() as conn:
        row = conn.execute(
            """
            SELECT 1
            FROM account_balance_snapshots
            WHERE account_id = ? AND source_kind = ? AND source_ref = ?
            """,
            (account_id, source_kind, source_ref),
        ).fetchone()
    return row is not None


def save_account_balance_snapshot(*, account_id: str, source_kind: str, source_ref: str = "", snapshot: dict) -> str:
    snapshot_id = str(uuid4())
    captured_at = utc_now_iso()
    with get_connection() as conn:
        conn.execute(
            """
            INSERT OR IGNORE INTO account_balance_snapshots
            (id, account_id, source_kind, source_ref, snapshot_json, captured_at)
            VALUES (?, ?, ?, ?, ?, ?)
            """,
            (snapshot_id, account_id, source_kind, source_ref, json.dumps(snapshot), captured_at),
        )
        row = conn.execute(
            """
            SELECT id
            FROM account_balance_snapshots
            WHERE account_id = ? AND source_kind = ? AND source_ref = ?
            """,
            (account_id, source_kind, source_ref),
        ).fetchone()
        conn.commit()
    if row is None:
        return snapshot_id
    return str(row["id"])


def list_recent_orders(*, account_id: str, client_id: str | None = None, status: str = "finished", limit: int = 50) -> list[dict]:
    query = """
        SELECT id, account_id, market, side, order_type, amount, price, expire_time, status,
               bitstamp_order_id, client_id, client_order_id, raw_json, created_at, updated_at
        FROM order_records
        WHERE account_id = ?
    """
    params: list[object] = [account_id]
    if client_id is not None:
        query += " AND client_id = ?"
        params.append(client_id)
    if status == "open":
        query += " AND lower(status) IN ('open', 'new', 'partially_filled')"
    elif status == "cancelled":
        query += " AND lower(status) = 'cancelled'"
    elif status == "finished":
        query += " AND lower(status) = 'finished'"
    elif status != "all":
        query += " AND lower(status) = ?"
        params.append(status)
    query += " ORDER BY created_at DESC, id DESC LIMIT ?"
    params.append(limit)

    with get_connection() as conn:
        rows = conn.execute(query, tuple(params)).fetchall()

    return [dict(row) for row in rows]


def list_account_balance_snapshots(*, account_id: str, limit: int = 50) -> list[dict]:
    with get_connection() as conn:
        rows = conn.execute(
            """
            SELECT id, account_id, source_kind, source_ref, snapshot_json, captured_at
            FROM account_balance_snapshots
            WHERE account_id = ?
            ORDER BY captured_at DESC, id DESC
            LIMIT ?
            """,
            (account_id, limit),
        ).fetchall()

    snapshots = []
    for row in rows:
        item = dict(row)
        item["snapshot"] = json.loads(item.pop("snapshot_json"))
        snapshots.append(item)
    return snapshots