import testutils/unittests
import ../ratelimit/store
import chronos
import db_connector/db_sqlite
import ../chat_sdk/migration
import std/[options, os, json]
import flatty

const dbName = "test_store.db"

suite "SqliteRateLimitStore Tests":
  setup:
    let db = open(dbName, "", "", "")
    runMigrations(db)

  teardown:
    if db != nil:
      db.close()
    if fileExists(dbName):
      removeFile(dbName)

  asyncTest "newSqliteRateLimitStore - empty state":
    ## Given
    let store = await RateLimitStore[string].new(db)

    ## When
    let loadedState = await store.loadBucketState()

    ## Then
    check loadedState.isNone()

  asyncTest "saveBucketState and loadBucketState - state persistence":
    ## Given
    let store = await RateLimitStore[string].new(db)

    let now = Moment.now()
    echo "now: ", now.epochSeconds()
    let newBucketState = BucketState(budget: 5, budgetCap: 20, lastTimeFull: now)

    ## When
    let saveResult = await store.saveBucketState(newBucketState)
    let loadedState = await store.loadBucketState()

    ## Then
    check saveResult == true
    check loadedState.isSome()
    check loadedState.get().budget == newBucketState.budget
    check loadedState.get().budgetCap == newBucketState.budgetCap
    check loadedState.get().lastTimeFull.epochSeconds() ==
      newBucketState.lastTimeFull.epochSeconds()

  asyncTest "queue operations - empty store":
    ## Given
    let store = await RateLimitStore[string].new(db)

    ## When/Then
    check store.getQueueLength(QueueType.Critical) == 0
    check store.getQueueLength(QueueType.Normal) == 0

    let criticalPop = await store.popFromQueue(QueueType.Critical)
    let normalPop = await store.popFromQueue(QueueType.Normal)

    check criticalPop.isNone()
    check normalPop.isNone()

  asyncTest "addToQueue and popFromQueue - single batch":
    ## Given
    let store = await RateLimitStore[string].new(db)
    let msgs = @[("msg1", "Hello"), ("msg2", "World")]

    ## When
    let addResult = await store.pushToQueue(QueueType.Critical, msgs)

    ## Then
    check addResult == true
    check store.getQueueLength(QueueType.Critical) == 1
    check store.getQueueLength(QueueType.Normal) == 0

    ## When
    let popResult = await store.popFromQueue(QueueType.Critical)

    ## Then
    check popResult.isSome()
    let poppedMsgs = popResult.get()
    check poppedMsgs.len == 2
    check poppedMsgs[0].msgId == "msg1"
    check poppedMsgs[0].msg == "Hello"
    check poppedMsgs[1].msgId == "msg2"
    check poppedMsgs[1].msg == "World"

    check store.getQueueLength(QueueType.Critical) == 0

  asyncTest "addToQueue and popFromQueue - multiple batches FIFO":
    ## Given
    let store = await RateLimitStore[string].new(db)
    let batch1 = @[("msg1", "First")]
    let batch2 = @[("msg2", "Second")]
    let batch3 = @[("msg3", "Third")]

    ## When - Add batches
    let result1 = await store.pushToQueue(QueueType.Normal, batch1)
    check result1 == true
    let result2 = await store.pushToQueue(QueueType.Normal, batch2)
    check result2 == true
    let result3 = await store.pushToQueue(QueueType.Normal, batch3)
    check result3 == true

    ## Then - Check lengths
    check store.getQueueLength(QueueType.Normal) == 3
    check store.getQueueLength(QueueType.Critical) == 0

    ## When/Then - Pop in FIFO order
    let pop1 = await store.popFromQueue(QueueType.Normal)
    check pop1.isSome()
    check pop1.get()[0].msg == "First"
    check store.getQueueLength(QueueType.Normal) == 2

    let pop2 = await store.popFromQueue(QueueType.Normal)
    check pop2.isSome()
    check pop2.get()[0].msg == "Second"
    check store.getQueueLength(QueueType.Normal) == 1

    let pop3 = await store.popFromQueue(QueueType.Normal)
    check pop3.isSome()
    check pop3.get()[0].msg == "Third"
    check store.getQueueLength(QueueType.Normal) == 0

    let pop4 = await store.popFromQueue(QueueType.Normal)
    check pop4.isNone()

  asyncTest "queue isolation - critical and normal queues are separate":
    ## Given
    let store = await RateLimitStore[string].new(db)
    let criticalMsgs = @[("crit1", "Critical Message")]
    let normalMsgs = @[("norm1", "Normal Message")]

    ## When
    let critResult = await store.pushToQueue(QueueType.Critical, criticalMsgs)
    check critResult == true
    let normResult = await store.pushToQueue(QueueType.Normal, normalMsgs)
    check normResult == true

    ## Then
    check store.getQueueLength(QueueType.Critical) == 1
    check store.getQueueLength(QueueType.Normal) == 1

    ## When - Pop from critical
    let criticalPop = await store.popFromQueue(QueueType.Critical)
    check criticalPop.isSome()
    check criticalPop.get()[0].msg == "Critical Message"

    ## Then - Normal queue unaffected
    check store.getQueueLength(QueueType.Critical) == 0
    check store.getQueueLength(QueueType.Normal) == 1

    ## When - Pop from normal
    let normalPop = await store.popFromQueue(QueueType.Normal)
    check normalPop.isSome()
    check normalPop.get()[0].msg == "Normal Message"

    ## Then - All queues empty
    check store.getQueueLength(QueueType.Critical) == 0
    check store.getQueueLength(QueueType.Normal) == 0

  asyncTest "queue persistence across store instances":
    ## Given
    let msgs = @[("persist1", "Persistent Message")]

    block:
      let store1 = await RateLimitStore[string].new(db)
      let addResult = await store1.pushToQueue(QueueType.Critical, msgs)
      check addResult == true
      check store1.getQueueLength(QueueType.Critical) == 1

    ## When - Create new store instance
    block:
      let store2 =await RateLimitStore[string].new(db)

      ## Then - Queue length should be restored from database
      check store2.getQueueLength(QueueType.Critical) == 1

      let popResult = await store2.popFromQueue(QueueType.Critical)
      check popResult.isSome()
      check popResult.get()[0].msg == "Persistent Message"
      check store2.getQueueLength(QueueType.Critical) == 0

  asyncTest "large batch handling":
    ## Given
    let store = await RateLimitStore[string].new(db)
    var largeBatch: seq[tuple[msgId: string, msg: string]]

    for i in 1 .. 100:
      largeBatch.add(("msg" & $i, "Message " & $i))

    ## When
    let addResult = await store.pushToQueue(QueueType.Normal, largeBatch)

    ## Then
    check addResult == true
    check store.getQueueLength(QueueType.Normal) == 1

    let popResult = await store.popFromQueue(QueueType.Normal)
    check popResult.isSome()
    let poppedMsgs = popResult.get()
    check poppedMsgs.len == 100
    check poppedMsgs[0].msgId == "msg1"
    check poppedMsgs[99].msgId == "msg100"
    check store.getQueueLength(QueueType.Normal) == 0