# Nim-Libp2p
# Copyright (c) 2023 Status Research & Development GmbH
# Licensed under either of
#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
#  * MIT license ([LICENSE-MIT](LICENSE-MIT))
# at your option.
# This file may not be copied, modified, or distributed except according to
# those terms.

when (NimMajor, NimMinor) < (1, 4):
  {.push raises: [Defect].}
else:
  {.push raises: [].}

## This module implements Pool of StreamTransport.
import chronos

const
  DefaultPoolSize* = 8
    ## Default pool size

type
  ConnectionFlags = enum
    None, Busy

  PoolItem = object
    transp*: StreamTransport
    flags*: set[ConnectionFlags]

  PoolState = enum
    Connecting, Connected, Closing, Closed

  TransportPool* = ref object
    ## Transports pool object
    transports: seq[PoolItem]
    busyCount: int
    state: PoolState
    bufferSize: int
    event: AsyncEvent

  TransportPoolError* = object of AsyncError

proc waitAll[T](futs: seq[Future[T]]): Future[void] =
  ## Performs waiting for all Future[T].
  var counter = len(futs)
  var retFuture = newFuture[void]("connpool.waitAllConnections")
  proc cb(udata: pointer) =
    dec(counter)
    if counter == 0:
      retFuture.complete()
  for fut in futs:
    fut.addCallback(cb)
  return retFuture

proc newPool*(address: TransportAddress, poolsize: int = DefaultPoolSize,
              bufferSize = DefaultStreamBufferSize,
             ): Future[TransportPool] {.async.} =
  ## Establish pool of connections to address ``address`` with size
  ## ``poolsize``.
  var pool = new TransportPool
  pool.bufferSize = bufferSize
  pool.transports = newSeq[PoolItem](poolsize)
  var conns = newSeq[Future[StreamTransport]](poolsize)
  pool.state = Connecting
  for i in 0..<poolsize:
    conns[i] = connect(address, bufferSize)
  # Waiting for all connections to be established.
  await waitAll(conns)
  # Checking connections and preparing pool.
  for i in 0..<poolsize:
    if conns[i].failed:
      raise conns[i].error
    else:
      let transp = conns[i].read()
      let item = PoolItem(transp: transp)
      pool.transports[i] = item
  # Setup available connections event
  pool.event = newAsyncEvent()
  pool.state = Connected
  result = pool

proc acquire*(pool: TransportPool): Future[StreamTransport] {.async.} =
  ## Acquire non-busy connection from pool ``pool``.
  var transp: StreamTransport
  if pool.state in {Connected}:
    while true:
      if pool.busyCount < len(pool.transports):
        for conn in pool.transports.mitems():
          if Busy notin conn.flags:
            conn.flags.incl(Busy)
            inc(pool.busyCount)
            transp = conn.transp
            break
      else:
        await pool.event.wait()
        pool.event.clear()

      if not isNil(transp):
        break
  else:
    raise newException(TransportPoolError, "Pool is not ready!")
  result = transp

proc release*(pool: TransportPool, transp: StreamTransport) =
  ## Release connection ``transp`` back to pool ``pool``.
  if pool.state in {Connected, Closing}:
    var found = false
    for conn in pool.transports.mitems():
      if conn.transp == transp:
        conn.flags.excl(Busy)
        dec(pool.busyCount)
        pool.event.fire()
        found = true
        break
    if not found:
      raise newException(TransportPoolError, "Transport not bound to pool!")
  else:
    raise newException(TransportPoolError, "Pool is not ready!")

proc join*(pool: TransportPool) {.async.} =
  ## Waiting for all connection to become available.
  if pool.state in {Connected, Closing}:
    while true:
      if pool.busyCount == 0:
        break
      else:
        await pool.event.wait()
        pool.event.clear()
  elif pool.state == Connecting:
    raise newException(TransportPoolError, "Pool is not ready!")

proc close*(pool: TransportPool) {.async.} =
  ## Closes transports pool ``pool`` and release all resources.
  if pool.state == Connected:
    pool.state = Closing
    # Waiting for all transports to become available.
    await pool.join()
    # Closing all transports
    var pending = newSeq[Future[void]](len(pool.transports))
    for i in 0..<len(pool.transports):
      let transp = pool.transports[i].transp
      transp.close()
      pending[i] = transp.join()
    # Waiting for all transports to be closed
    await waitAll(pending)
    # Mark pool as `Closed`.
    pool.state = Closed