//// A worker pool for Gleam that distributes work across a limited number of
//// worker processes. Workers are pooled to avoid the overhead of
//// spawning and killing processes for each task.
////
//// The pool manages a queue of work items and distributes them to idle workers.
//// When no workers are available, work is queued until a worker becomes free.
//// The pool handles worker crashes gracefully and automatically manages the
//// worker lifecycle.
////
//// ## Example
////
//// ```gleam
//// import crew
//// import gleam/otp/static_supervisor as supervisor
////
//// pub fn main() {
////   // Create a pool name
////   let pool_name = process.new_name("my_crew")
////
////   // Start an unsupervised pool
////   let assert Ok(_) =
////       crew.new(pool_name)
////       |> crew.fixed_size(4)
////       |> crew.start
////
////   // Execute work on the pool
////   let result = crew.work(pool_name, 5000, fn() {
////     // Some expensive computation
////     expensive_computation()
////   })
//// }
//// ```

// -- IMPORTS -----------------------------------------------------------------

import gleam/bool
import gleam/deque.{type Deque}
import gleam/dict.{type Dict}
import gleam/erlang/process.{
  type Monitor, type Name, type Pid, type Selector, type Subject,
}
import gleam/int
import gleam/list
import gleam/otp/actor.{type Next, type StartError}
import gleam/otp/static_supervisor.{type Supervisor}
import gleam/otp/supervision.{type ChildSpecification}
import gleam/result
import gleam/set.{type Set}
import gleam/string

// TODO: - min/max dynamic scaling? what (if any) is the load metric?
// TODO: fairer scheduler? group work by caller and round-robin maybe?

// poolboy: https://github.com/devinus/poolboy
// lifeguard: https://github.com/Pevensie/lifeguard
// gen_stage: https://github.com/elixir-lang/gen_stage

// -- BUILDER -----------------------------------------------------------------

/// A builder for configuring a worker pool before starting it.
pub opaque type Builder {
  Builder(name: Name(PoolMsg), size: Int)
}

/// Create a new worker pool builder with the given name.
///
/// The name is used to register the pool so that work can be sent to it.
///
/// By default, the pool will have a number of workers equal to the number of
/// scheduler threads available on the system (typically the number of CPU cores).
///
/// ## Example
///
/// ```gleam
/// let pool_name = process.new_name("pool")
/// let builder = crew.new(pool_name)
/// ```
pub fn new(name: Name(PoolMsg)) -> Builder {
  Builder(name:, size: scheduler_count())
}

/// Set the number of worker processes in the pool to a fixed number.
///
/// When set to less than 1 starting the pool will fail.
///
/// ## Example
///
/// ```gleam
/// crew.new(pool_name)
/// |> crew.fixed_size(8)  // Use 8 workers regardless of CPU count
/// ```
pub fn fixed_size(builder: Builder, size: Int) -> Builder {
  Builder(..builder, size:)
}

// -- START -------------------------------------------------------------------

/// Start an unsupervised worker pool from the given builder.
///
/// Returns a supervisor that manages the pool and its workers. In most cases,
/// you should use `supervised` instead to get a child specification that can
/// be added to your application's supervision tree.
///
/// ## Panics
/// This function will exit the process if any workers fail to start, similar
/// to `static_supervisor.start`.
///
/// ## Example
///
/// ```gleam
/// let assert Ok(pool_supervisor) =
///   crew.new(pool_name)
///   |> crew.start
/// ```
pub fn start(builder: Builder) -> Result(Supervisor, StartError) {
  let Builder(name:, size:) = builder
  use result <- result.map(start_tree(name, size))
  result.data
}

/// Create a child specification for a supervised worker pool.
///
/// This is the recommended way to start a worker pool as part of your
/// application's supervision tree. The returned child specification can be
/// added to a supervisor using `static_supervisor.add`.
///
/// ## Example
///
/// ```gleam
/// let pool_spec =
///   crew.new(pool_name)
///   |> crew.fixed_size(4)
///   |> crew.supervised
///
/// let assert Ok(_) =
///   supervisor.new(supervisor.OneForOne)
///   |> supervisor.add(pool_spec)
///   |> supervisor.start
/// ```
pub fn supervised(builder: Builder) -> ChildSpecification(_) {
  let Builder(name:, size:) = builder
  use <- supervision.supervisor
  start_tree(name, size)
}

fn start_tree(
  name: Name(PoolMsg),
  size: Int,
) -> Result(actor.Started(_), StartError) {
  use <- bool.guard(
    when: size <= 0,
    return: Error(actor.InitFailed("pool size must be greater than zero")),
  )

  let main_supervisor = static_supervisor.new(static_supervisor.RestForOne)
  let worker_supervisor = static_supervisor.new(static_supervisor.OneForOne)

  let pool_spec = {
    use <- supervision.worker

    actor.new_with_initialiser(1000, init_pool)
    |> actor.named(name)
    |> actor.on_message(pool)
    |> actor.start
  }

  let worker_spec = {
    use <- supervision.worker
    let pid = process.spawn(fn() { worker(name) })
    Ok(actor.Started(pid, Nil))
  }

  let worker_supervisor_spec = {
    use <- supervision.supervisor

    worker_supervisor
    |> repeat(times: size, with: static_supervisor.add(_, worker_spec))
    |> static_supervisor.start
  }

  main_supervisor
  |> static_supervisor.add(pool_spec)
  |> static_supervisor.add(worker_supervisor_spec)
  |> static_supervisor.start()
}

// -- WORK --------------------------------------------------------------------

/// Execute a single piece of work on the pool and wait for the result.
///
/// This function blocks until the work is completed or the timeout is reached.
/// The work function is executed on one of the worker processes in the pool.
///
/// ## Parameters
/// - `pool` - The name of the pool to execute work on
/// - `timeout` - Maximum time to wait for completion in milliseconds
/// - `work` - A function containing the work to be executed
///
/// ## Panics
/// - If the pool does not complete the work within the specified timeout
/// - If the pool is not running
/// - If the worker crashes while executing the work
///
/// ## Example
///
/// ```gleam
/// let response = crew.work(pool_name, 5000, fn() {
///   httpc.send(...)
/// })
/// ```
pub fn work(
  in pool: Name(PoolMsg),
  timeout timeout: Int,
  do work: fn() -> any,
) -> any {
  let assert [result] = work_many(pool, timeout, [work])
  result
}

/// Execute multiple pieces of work concurrently on the pool, without
/// ordering guarantees.
///
/// Work is distributed among the available workers and executed concurrently.
/// Results are returned in the order they complete, not the order they were
/// submitted.
///
/// ## Parameters
/// - `pool` - The name of the pool to execute work on
/// - `timeout` - Maximum time to wait for all work to complete in milliseconds
/// - `work` - A list of functions containing work to be executed
///
/// ## Panics
/// - If the pool does not complete all work within the specified timeout
/// - If the pool is not running
/// - If any worker crashes while executing work
///
/// ## Example
///
/// ```gleam
/// let users = crew.work_many(pool_name, 10000, [
///   fn() { fetch_user_data(user1) },
///   fn() { fetch_user_data(user2) },
///   fn() { fetch_user_data(user3) },
/// ])
/// ```
pub fn work_many(
  in pool: Name(PoolMsg),
  timeout timeout: Int,
  do work: List(fn() -> any),
) -> List(any) {
  list.reverse(do_work(pool, timeout, work))
}

/// Execute multiple pieces of work concurrently on the pool and return results
/// in submission order.
///
/// Work functions are distributed across available workers and executed
/// concurrently, but results are reordered to match the original submission
/// order before being returned.
///
/// ## Parameters
/// - `pool` - The name of the pool to execute work on
/// - `timeout` - Maximum time to wait for all work to complete in milliseconds
/// - `work` - A list of functions containing work to be executed
///
/// ## Panics
/// - If the pool does not complete all work within the specified timeout
/// - If the pool is not running
/// - If any worker crashes while executing work
///
/// ## Example
///
/// ```gleam
/// let assert [user1, user2, user3] = crew.work_ordered(pool_name, 10000, [
///   fn() { fetch_user_data(user1) },
///   fn() { fetch_user_data(user2) },
///   fn() { fetch_user_data(user3) },
/// ])
/// ```
pub fn work_ordered(
  in pool: Name(PoolMsg),
  timeout timeout: Int,
  do work: List(fn() -> any),
) -> List(any) {
  parallel_map(work, pool, timeout, fn(f) { f() })
}

/// Apply a function to each element of a list concurrently using the worker pool.
///
/// This is similar to `list.map` but executes the mapping function
/// concurrently across all workers. Results are returned in the same
/// order as the input list.
///
/// There's a bunch of extra overhead involved with spawning a work item
/// per list element and making sure the order matches. Depending on your
/// workload it might make sense to split your list into chunks first to reduce
/// work queue pressure.
///
/// ## Parameters
/// - `list` - The list of items to map over
/// - `pool` - The name of the pool to execute work on
/// - `timeout` - Maximum time to wait for all mappings to complete in milliseconds
/// - `fun` - The function to apply to each item
///
/// ## Panics
/// - If the pool does not complete all mappings within the specified timeout
/// - If the pool is not running
/// - If any worker crashes while executing a mapping
///
/// ## Example
///
/// ```gleam
/// let user_ids = [1, 2, 3, 4, 5]
/// let users = crew.parallel_map(user_ids, pool_name, 5000, fetch_user_data)
/// ```
pub fn parallel_map(
  over list: List(a),
  in pool: Name(PoolMsg),
  timeout timeout: Int,
  with fun: fn(a) -> b,
) -> List(b) {
  let unordered: List(#(Int, b)) =
    list
    |> list.index_map(fn(item, index) { fn() { #(index, fun(item)) } })
    |> do_work(pool, timeout, _)

  unordered
  |> list.sort(fn(a, b) { int.compare(a.0, b.0) })
  |> list.map(fn(x) { x.1 })
}

fn do_work(
  pool: Name(PoolMsg),
  timeout: Int,
  work: List(fn() -> any),
) -> List(any) {
  let channel = subscribe(pool)

  enqueue_many(channel, work)

  let selector =
    process.new_selector()
    |> select_map(channel, fn(x) { x })

  let timeout_end = system_time() + timeout
  let result = receive_loop(work, selector, timeout_end, [])

  // sending unsubscribe unconditionally in an attempt to handle
  // cases more gracefully where this crash is rescued by the caller.
  unsubscribe(channel, selector)

  case result {
    Ok(value) -> value
    Error(_) -> panic as "Pool did not complete work in the alloted time"
  }
}

fn receive_loop(
  work: List(fn() -> any),
  selector: Selector(any),
  timeout_end: Int,
  state: List(any),
) -> Result(List(any), Nil) {
  case work {
    [] -> Ok(state)
    [_, ..work] -> {
      let timeout = timeout_end - system_time()
      case process.selector_receive(selector, timeout) {
        Ok(result) ->
          receive_loop(work, selector, timeout_end, [result, ..state])
        Error(Nil) -> Error(Nil)
      }
    }
  }
}

// -- SUBSCRIBE / UNSUBSCRIBE -------------------------------------------------

/// A channel represents a typed subscription to a worker pool for receiving
/// work results. This allows you to send work and receive results asynchronously.
pub opaque type Channel(a) {
  Channel(
    pool: Subject(PoolMsg),
    receive: Subject(WorkResult),
    monitor: Monitor,
  )
}

/// Create a typed subscription channel to the worker pool for work submission.
///
/// This function provides a lower-level interface for submitting work to the
/// pool. In many cases calling the `work*` function from a separate process
/// will be easier.
///
/// Subscribing to the pool creates a reference to the current process in the pool
/// that must be cleaned up using `unsubscribe`. The returned channel can be used
/// with `select_map` and `enqueue` for custom message handling.
///
/// ## Panics
/// - If the pool is not running
///
/// ## Example
///
/// ```gleam
/// let channel = crew.subscribe(pool_name)
/// crew.enqueue(channel, fn() { some_work() })
/// // Handle results with selector...
/// crew.unsubscribe(channel, selector)
/// ```
pub fn subscribe(pool: Name(PoolMsg)) -> Channel(a) {
  let pool_subject = process.named_subject(pool)
  let assert Ok(pool_pid) = process.subject_owner(pool_subject)
    as "Pool is not running"

  let monitor = process.monitor(pool_pid)

  let receive: Subject(WorkResult) = process.new_subject()
  actor.send(pool_subject, Subscribe(receive))

  Channel(pool: pool_subject, receive:, monitor:)
}

/// Add a worker pool channel to a selector for receiving work results.
///
/// This allows you to receive work results as part of the larger message-handling
/// loop. Used in conjunction with `subscribe` and `enqueue` for lower-level
/// pool usage.
///
/// Work results arrive in completion order. It is your responsibility to handle
/// the lifecycle and ordering.
///
/// ## Panics
/// - If a worker or the pool crashes while executing work, the selector will
/// panic with details about the crash.
pub fn select_map(
  selector: Selector(msg),
  channel: Channel(a),
  tagger: fn(a) -> msg,
) -> Selector(msg) {
  selector
  |> process.select_specific_monitor(channel.monitor, fn(down) {
    let msg =
      "Pool exited while waiting for work to complete: " <> string.inspect(down)

    panic as msg
  })
  |> process.select_map(channel.receive, fn(result) {
    case result {
      Done(value) -> tagger(cast(value))
      WorkerExited(reason) -> {
        let msg = "Worker exited: " <> string.inspect(reason)
        panic as msg
      }
    }
  })
}

/// Remove a subscription channel from the pool and clean up resources.
///
/// This function should be called when you're done using a channel obtained
/// from `subscribe`. It stops any in-progress work for this channel and
/// removes the channel's handlers from the selector.
pub fn unsubscribe(
  channel: Channel(a),
  selector: Selector(msg),
) -> Selector(msg) {
  actor.send(channel.pool, Unsubscribe(channel.receive))

  process.demonitor_process(channel.monitor)

  selector
  |> process.deselect_specific_monitor(channel.monitor)
  |> process.deselect(channel.receive)
}

/// Submit a single piece of work to the pool using a subscription channel.
///
/// This is a lower-level function for submitting work. Results must be handled
/// using a selector with `select_map`. It is the callers responsibility to handle
/// timeouts and submission order. Most users should prefer the `work*` functions.
pub fn enqueue(channel: Channel(a), work: fn() -> a) -> Nil {
  enqueue_many(channel, [work])
}

/// Submit multiple pieces of work to the pool using a subscription channel.
///
/// This is a lower-level function for submitting work. Results must be handled
/// using a selector with `select_map`. It is the callers responsibility to handle
/// timeouts and submission order. Most users should prefer the `work*` functions.
pub fn enqueue_many(channel: Channel(a), work: List(fn() -> a)) -> Nil {
  actor.send(channel.pool, Enqueue(channel.receive, cast(work)))
}

// -- POOL --------------------------------------------------------------------

pub opaque type PoolMsg {
  WorkerStarted(pid: Pid, send: Subject(Work))
  WorkerIdle(pid: Pid)
  MonitoredProcessExited(reason: process.Down)
  //
  Subscribe(receive: Subject(WorkResult))
  Unsubscribe(receive: Subject(WorkResult))
  Enqueue(receive: Subject(WorkResult), work: List(fn() -> Any))
}

type State {
  State(
    idle_workers: List(Worker),
    active_workers: Dict(Pid, ActiveWorker),
    callers: Dict(Pid, Caller),
    requests: Dict(Subject(WorkResult), Request),
    queue: Deque(QueueItem),
  )
}

type Worker {
  Worker(pid: Pid, send: Subject(Work), monitor: Monitor)
}

type QueueItem {
  QueueItem(work: fn() -> Any, caller: Pid, receive: Subject(WorkResult))
}

type ActiveWorker {
  ActiveWorker(worker: Worker, caller: Pid, receive: Subject(WorkResult))
}

type Caller {
  Caller(pid: Pid, monitor: Monitor, requests: Set(Subject(WorkResult)))
}

type Request {
  Request(from: Pid, workers: Set(Pid), receive: Subject(WorkResult))
}

fn init_pool(self: Subject(PoolMsg)) {
  let selector =
    process.new_selector()
    |> process.select(self)
    |> process.select_monitors(MonitoredProcessExited)

  let state =
    State(
      idle_workers: [],
      active_workers: dict.new(),
      callers: dict.new(),
      requests: dict.new(),
      queue: deque.new(),
    )

  actor.initialised(state)
  |> actor.selecting(selector)
  |> actor.returning(self)
  |> Ok
}

fn pool(state: State, msg: PoolMsg) -> Next(State, PoolMsg) {
  let next_state = case msg {
    WorkerStarted(pid:, send:) -> {
      let monitor = process.monitor(pid)

      let worker = Worker(pid:, send:, monitor:)
      try_dequeue_work(state, worker)
    }

    MonitoredProcessExited(process.ProcessDown(pid:, monitor: _, reason:)) -> {
      case dict.get(state.callers, pid), dict.get(state.active_workers, pid) {
        // Caller down -
        // We will get a second Down message for when the process is caller and worker.
        Ok(_), _ -> {
          // this sends a message back to the caller, which we already
          // know is down. This should not matter.
          abort_caller(state, pid, reason)
        }

        // Active worker down
        Error(_), Ok(ActiveWorker(caller:, ..)) -> {
          // after one worker exits, we want to stop all workers for this caller
          // and remove it as an active caller. The caller will no longer receive messages.
          let active_workers = dict.delete(state.active_workers, pid)

          State(..state, active_workers:)
          |> abort_caller(caller, reason)
        }
        // idle worker down
        Error(_), Error(_) -> {
          let idle_workers =
            list.filter(state.idle_workers, fn(worker) { worker.pid != pid })

          State(..state, idle_workers:)
        }
      }
    }

    // we only monitor processes and can ignore all other Down messages
    MonitoredProcessExited(_) -> state

    WorkerIdle(pid:) -> {
      // got an idle message from a worker that is not active
      use ActiveWorker(worker:, receive:, caller: _) <- try_(
        dict.get(state.active_workers, pid),
        state,
      )

      let active_workers = dict.delete(state.active_workers, pid)

      case dict.get(state.requests, receive) {
        Ok(request) -> {
          let request =
            Request(..request, workers: set.delete(request.workers, pid))
          let requests = dict.insert(state.requests, receive, request)

          State(..state, requests:, active_workers:)
          |> try_dequeue_work(worker)
        }

        // the requests unsubscribed while we were still working.
        // we sent a kill request while unsubscribing, so we drop the worker here.
        Error(_) -> State(..state, active_workers:)
      }
    }

    Enqueue(work: [], ..) -> state
    Enqueue(receive:, work:) -> {
      use request <- try_(dict.get(state.requests, receive), state)
      enqueue_loop(state, request, work)
    }

    Subscribe(receive:) -> {
      use pid <- try_(process.subject_owner(receive), state)

      let request =
        dict.get(state.requests, receive)
        |> result.unwrap(Request(from: pid, workers: set.new(), receive:))

      let caller = case dict.get(state.callers, pid) {
        Ok(caller) ->
          Caller(..caller, requests: set.insert(caller.requests, receive))

        Error(_) -> {
          let monitor = process.monitor(pid)
          Caller(pid: pid, monitor:, requests: set.new() |> set.insert(receive))
        }
      }

      let callers = dict.insert(state.callers, pid, caller)
      let requests = dict.insert(state.requests, receive, request)

      State(..state, callers:, requests:)
    }

    Unsubscribe(receive:) -> {
      use request <- try_(dict.get(state.requests, receive), state)

      let requests = dict.delete(state.requests, receive)

      set.each(request.workers, process.kill)

      use caller <- try_(
        dict.get(state.callers, request.from),
        State(..state, requests:),
      )

      let caller =
        Caller(..caller, requests: set.delete(caller.requests, receive))
      let callers = case set.is_empty(caller.requests) {
        True -> {
          process.demonitor_process(caller.monitor)
          dict.delete(state.callers, caller.pid)
        }
        False -> {
          dict.insert(state.callers, caller.pid, caller)
        }
      }

      State(..state, requests:, callers:)
    }
  }

  actor.continue(next_state)
}

fn enqueue_loop(
  state: State,
  request: Request,
  work: List(fn() -> Any),
) -> State {
  let Request(from: caller, receive:, workers:) = request

  case work, state.idle_workers {
    [work, ..rest], [worker, ..idle_workers] -> {
      // got a work item and a worker - it's a match!
      process.send(worker.send, Work(work:, receive:))

      let active_worker = ActiveWorker(worker:, caller:, receive:)
      let active_workers =
        dict.insert(state.active_workers, worker.pid, active_worker)

      let request = Request(..request, workers: set.insert(workers, worker.pid))

      State(..state, active_workers:, idle_workers:)
      |> enqueue_loop(request, rest)
    }

    _, _ -> {
      let queue =
        list.fold(work, state.queue, fn(queue, work) {
          deque.push_back(queue, QueueItem(work:, caller:, receive:))
        })

      let requests = dict.insert(state.requests, receive, request)

      State(..state, queue:, requests:)
    }
  }
}

fn try_dequeue_work(state: State, worker: Worker) -> State {
  use #(QueueItem(work:, caller:, receive:), queue) <- try(
    deque.pop_front(state.queue),
    // the queue is empty, this worker becomes idle.
    fn(_) { State(..state, idle_workers: [worker, ..state.idle_workers]) },
  )

  use request <- try(dict.get(state.requests, receive), fn(_) {
    // this request unsubscribed
    try_dequeue_work(State(..state, queue:), worker)
  })

  process.send(worker.send, Work(work:, receive:))

  let active_worker = ActiveWorker(worker:, caller:, receive:)
  let active_workers =
    dict.insert(state.active_workers, worker.pid, active_worker)

  let request =
    Request(..request, workers: set.insert(request.workers, worker.pid))
  let requests = dict.insert(state.requests, receive, request)

  State(..state, active_workers:, requests:, queue:)
}

fn abort_caller(state: State, caller: Pid, reason: process.ExitReason) -> State {
  // we do not have an active caller, so this worker was stopped in
  // response to a caller exiting.
  use caller <- try_(dict.get(state.callers, caller), state)

  process.demonitor_process(caller.monitor)

  let requests =
    set.fold(caller.requests, state.requests, fn(requests, receive) {
      use request <- try_(dict.get(state.requests, receive), requests)
      // we kill all workers currently working on things for this caller.
      // the WorkerStopped messages will then clean up the workers.
      set.each(request.workers, process.kill)

      // notify the caller that we crashed. this is the last message we send.
      process.send(receive, WorkerExited(reason))

      dict.delete(requests, receive)
    })

  // the queue is cleaned up while consuming it if the caller is missing.
  let callers = dict.delete(state.callers, caller.pid)

  State(..state, callers:, requests:)
}

// -- WORKER ------------------------------------------------------------------

type Any

type Work {
  Work(work: fn() -> Any, receive: Subject(WorkResult))
}

type WorkResult {
  Done(Any)
  WorkerExited(process.ExitReason)
}

fn worker(pool: Name(PoolMsg)) -> Nil {
  let self = process.self()
  let subject = process.new_subject()
  let pool_subject = process.named_subject(pool)

  process.send(pool_subject, WorkerStarted(self, subject))
  worker_loop(pool_subject, self, subject)
}

fn worker_loop(pool: Subject(PoolMsg), self: Pid, subject: Subject(Work)) {
  let Work(work:, receive:) = process.receive_forever(subject)

  let result = work()
  process.send(pool, WorkerIdle(self))
  process.send(receive, Done(result))

  worker_loop(pool, self, subject)
}

// -- HELPERS -----------------------------------------------------------------

fn repeat(
  times times: Int,
  from state: state,
  with f: fn(state) -> state,
) -> state {
  case times > 0 {
    True -> repeat(times - 1, f(state), f)
    False -> state
  }
}

fn try(result: Result(a, x), or: fn(x) -> b, then: fn(a) -> b) {
  case result {
    Ok(x) -> then(x)
    Error(x) -> or(x)
  }
}

fn try_(result: Result(a, _), or: b, then: fn(a) -> b) {
  case result {
    Ok(x) -> then(x)
    Error(_) -> or
  }
}

@external(erlang, "crew_ffi", "scheduler_count")
fn scheduler_count() -> Int

@external(erlang, "crew_ffi", "system_time")
fn system_time() -> Int

@external(erlang, "crew_ffi", "identity")
fn cast(value: a) -> b