nim-chat-sdk/ratelimit/store/sqlite.nim

import std/times
import std/strutils
import std/json
import ./store
import chronos
import db_connector/db_sqlite

# SQLite Implementation
type SqliteRateLimitStore* = ref object
  db: DbConn
  dbPath: string

const BUCKET_STATE_KEY = "rate_limit_bucket_state"

proc newSqliteRateLimitStore*(dbPath: string = ":memory:"): SqliteRateLimitStore =
  result = SqliteRateLimitStore(dbPath: dbPath)
  result.db = open(dbPath, "", "", "")

  # Create table if it doesn't exist
  result.db.exec(
    sql"""
    CREATE TABLE IF NOT EXISTS kv_store (
        key TEXT PRIMARY KEY,
        value BLOB
    )
  """
  )

  # Insert default state if key doesn't exist
  let count = result.db
    .getValue(sql"SELECT COUNT(*) FROM kv_store WHERE key = ?", BUCKET_STATE_KEY)
    .parseInt()
  if count == 0:
    let defaultTimeSeconds = Moment.now().epochSeconds()
    let defaultState = BucketState(
      budget: 10,
      budgetCap: 10,
      lastTimeFull: Moment.init(defaultTimeSeconds, chronos.seconds(1)),
    )

    # Serialize to JSON
    let jsonState =
      %*{
        "budget": defaultState.budget,
        "budgetCap": defaultState.budgetCap,
        "lastTimeFullSeconds": defaultTimeSeconds,
      }

    result.db.exec(
      sql"INSERT INTO kv_store (key, value) VALUES (?, ?)", BUCKET_STATE_KEY, $jsonState
    )

proc close*(store: SqliteRateLimitStore) =
  if store.db != nil:
    store.db.close()

proc saveBucketState*(
    store: SqliteRateLimitStore, bucketState: BucketState
): Future[bool] {.async.} =
  try:
    # Convert Moment to Unix seconds for storage
    let lastTimeSeconds = bucketState.lastTimeFull.epochSeconds()

    # Serialize to JSON
    let jsonState =
      %*{
        "budget": bucketState.budget,
        "budgetCap": bucketState.budgetCap,
        "lastTimeFullSeconds": lastTimeSeconds,
      }

    store.db.exec(
      sql"UPDATE kv_store SET value = ? WHERE key = ?", $jsonState, BUCKET_STATE_KEY
    )
    return true
  except:
    return false

proc loadBucketState*(store: SqliteRateLimitStore): Future[BucketState] {.async.} =
  let jsonStr =
    store.db.getValue(sql"SELECT value FROM kv_store WHERE key = ?", BUCKET_STATE_KEY)

  # Parse JSON and reconstruct BucketState
  let jsonData = parseJson(jsonStr)
  let unixSeconds = jsonData["lastTimeFullSeconds"].getInt().int64
  let lastTimeFull = Moment.init(unixSeconds, chronos.seconds(1))

  return BucketState(
    budget: jsonData["budget"].getInt(),
    budgetCap: jsonData["budgetCap"].getInt(),
    lastTimeFull: lastTimeFull,
  )