import pkg/chronos
import pkg/asynctest
import pkg/stew/results

import pkg/codex/utils/asyncheapqueue
import pkg/codex/rng

type
  Task* = tuple[name: string, priority: int]

proc `<`*(a, b: Task): bool =
  a.priority < b.priority

proc `==`*(a, b: Task): bool =
  a.name == b.name

proc toSortedSeq[T](h: AsyncHeapQueue[T], queueType = QueueType.Min): seq[T] =
  var tmp = newAsyncHeapQueue[T](queueType = queueType)
  for d in h:
    check tmp.pushNoWait(d).isOk
  while tmp.len > 0:
    result.add(popNoWait(tmp).tryGet())

suite "Synchronous tests":
  test "Test pushNoWait - Min":
    var heap = newAsyncHeapQueue[int]()
    let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
    for item in data:
      check heap.pushNoWait(item).isOk

    check heap[0] == 0
    check heap.toSortedSeq == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

  test "Test pushNoWait - Max":
    var heap = newAsyncHeapQueue[int](queueType = QueueType.Max)
    let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
    for item in data:
      check heap.pushNoWait(item).isOk

    check heap[0] == 9
    check heap.toSortedSeq(QueueType.Max) == @[9, 8, 7, 6, 5, 4, 3, 2, 1, 0]

  test "Test popNoWait":
    var heap = newAsyncHeapQueue[int]()
    let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
    for item in data:
      check heap.pushNoWait(item).isOk

    var res: seq[int]
    while heap.len > 0:
      let r = heap.popNoWait()
      if r.isOk:
        res.add(r.get)

    check res == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

  test "Test popNoWait - Max":
    var heap = newAsyncHeapQueue[int](queueType = QueueType.Max)
    let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
    for item in data:
      check heap.pushNoWait(item).isOk

    var res: seq[int]
    while heap.len > 0:
      let r = heap.popNoWait()
      if r.isOk:
        res.add(r.get)

    check res == @[9, 8, 7, 6, 5, 4, 3, 2, 1, 0]

  test "Test del": # Test del
    var heap = newAsyncHeapQueue[int]()
    let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
    for item in data:
      check heap.pushNoWait(item).isOk

    heap.del(0)
    doAssert(heap[0] == 1)

    heap.del(heap.find(7))
    check heap.toSortedSeq == @[1, 2, 3, 4, 5, 6, 8, 9]

    heap.del(heap.find(5))
    check heap.toSortedSeq == @[1, 2, 3, 4, 6, 8, 9]

    heap.del(heap.find(6))
    check heap.toSortedSeq == @[1, 2, 3, 4, 8, 9]

    heap.del(heap.find(2))
    check heap.toSortedSeq == @[1, 3, 4, 8, 9]

  test "Test del last": # Test del last
    var heap = newAsyncHeapQueue[int]()
    let data = [1, 2, 3]
    for item in data:
      check heap.pushNoWait(item).isOk

    heap.del(2)
    check heap.toSortedSeq == @[1, 2]

    heap.del(1)
    check heap.toSortedSeq == @[1]

    heap.del(0)
    check heap.toSortedSeq == newSeq[int]() # empty seq has no type

  test "Should throw popping from an empty queue":
    var heap = newAsyncHeapQueue[int]()
    let err = heap.popNoWait()
    check err.isErr
    check err.error == AsyncHQErrors.Empty

  test "Should throw pushing to an full queue":
    var heap = newAsyncHeapQueue[int](1)
    check heap.pushNoWait(1).isOk
    let err = heap.pushNoWait(2)
    check err.isErr
    check err.error == AsyncHQErrors.Full

  test "Test clear":
    var heap = newAsyncHeapQueue[int]()
    let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
    for item in data:
      check heap.pushNoWait(item).isOk

    check heap.len == 10
    heap.clear()
    check heap.len == 0

suite "Asynchronous Tests":

  test "Test push":
    var heap = newAsyncHeapQueue[int]()
    let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
    for item in data:
      await push(heap, item)
    check heap[0] == 0
    check heap.toSortedSeq == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

  test "Test push and pop with maxSize":
    var heap = newAsyncHeapQueue[int](5)
    let data = [1, 9, 5, 3, 7, 4, 2]

    proc pushTask() {.async.} =
      for item in data:
        await push(heap, item)

    asyncSpawn pushTask()

    check heap.len == 5
    check heap[0] == 1 # because we haven't pushed 0 yet

    check (await heap.pop) == 1
    check (await heap.pop) == 3
    check (await heap.pop) == 5
    check (await heap.pop) == 7
    check (await heap.pop) == 9

    await sleepAsync(1.milliseconds) # allow poll to run once more
    check (await heap.pop) == 2
    check (await heap.pop) == 4

  test "Test update":
    var heap = newAsyncHeapQueue[Task](5)

    for item in [("a", 4), ("b", 3), ("c", 2)]:
      check heap.pushNoWait(item).isOk

    check heap[0] == (name: "c", priority: 2)
    check heap.update((name: "a", priority: 1))
    check heap[0] == (name: "a", priority: 1)

  test "Test pushOrUpdate - update":
    var heap = newAsyncHeapQueue[Task](3)

    for item in [("a", 4), ("b", 3), ("c", 2)]:
      check heap.pushNoWait(item).isOk

    check heap[0] == (name: "c", priority: 2)
    await heap.pushOrUpdate((name: "a", priority: 1))
    check heap[0] == (name: "a", priority: 1)

  test "Test pushOrUpdate - push":
    var heap = newAsyncHeapQueue[Task](2)

    for item in [("a", 4), ("b", 3)]:
      check heap.pushNoWait(item).isOk

    check heap[0] == ("b", 3)                 # sanity check for order

    let fut = heap.pushOrUpdate(("c", 2))     # attempt to push a non existen item but block
    check heap.popNoWait().tryGet() == ("b", 3)  # pop one off
    await fut                                 # wait for push to complete

    check heap[0] == (name: "c", priority: 2) # check order again

  test "Test pop":
    var heap = newAsyncHeapQueue[int]()
    let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
    for item in data:
      check heap.pushNoWait(item).isOk

    var res: seq[int]
    while heap.len > 0:
      res.add((await heap.pop()))

    check res == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

  test "Test delete":
    var heap = newAsyncHeapQueue[Task]()
    let data = ["d", "b", "c", "a", "h", "e", "f", "g"]

    for item in data:
      check heap.pushNoWait((
        name: item,
        priority: Rng.instance().rand(data.len)
      )).isOk

    let del = heap[3]
    heap.delete(del)
    check heap.find(del) < 0