// Connection establishment and top-level client lifecycle APIs.

///|
/// Construct a client handle and its single-use executor without network I/O.
/// Spawn `executor.run()` and await `client.ready()` before reading startup data.
/// The asynchronous message buffer holds 256 entries by default.
pub fn Client::create(config : Config) -> (Client, ClientExecutor) {
  Client::create_with_capacity_unchecked(config, 256)
}

///|
/// Create a client with a bounded asynchronous message buffer.
/// The oldest message is discarded when the buffer is full.
pub fn Client::create_with_async_message_capacity(
  config : Config,
  capacity : Int32,
) -> (Client, ClientExecutor) raise {
  if capacity < 1 {
    raise ClientError::InvalidConfig("async message capacity must be positive")
  }
  Client::create_with_capacity_unchecked(config, capacity)
}

///|
fn Client::create_with_capacity_unchecked(
  config : Config,
  capacity : Int32,
) -> (Client, ClientExecutor) {
  let types : Map[@proto.Oid, Type] = Map([])
  seed_builtin_types(types)
  let shared = {
    config,
    requests: Queue(kind=Unbounded),
    async_messages: Queue(kind=Blocking(capacity)),
    dropped_async_messages: Ref(0L),
    async_message_handler: Ref(None),
    connection_id: connection_ids.allocate(),
    root_transaction_occupancy: Ref(0),
    control: DrainControl::new(),
    ready: Completion::new(),
    done: Completion::new(),
    started: Ref(false),
    runtime_error: Ref(None),
    parameters: Map([]),
    types,
    process_id: Ref(0),
    secret_key: Ref(0),
    transaction_status: Ref(b'I'),
    closed: Ref(false),
    closing: Ref(false),
    close_requested: Ref(false),
    aborted: Ref(false),
    operation_gate: OperationGate::new(),
    driver_task: Ref(None),
    next_id: Ref(0),
  }
  ({ shared, }, { shared, })
}

///|
/// Number of asynchronous messages discarded from this client's own buffer.
pub fn Client::dropped_async_messages(self : Client) -> Int64 {
  self.shared.dropped_async_messages.val
}

///|
/// Route asynchronous messages to the pool before `run`.
/// Installed handlers receive messages instead of the `next_message()` buffer.
#internal(internal, "for use by pgpool only")
#doc(hidden)
pub fn Client::install_async_message_handler(
  self : Client,
  handler : (AsyncMessage) -> Unit,
) -> Unit raise {
  if self.shared.started.val || self.shared.async_message_handler.val is Some(_) {
    raise ClientError::Protocol(
      "async message handler must be installed before executor start",
    )
  }
  self.shared.async_message_handler.val = Some(handler)
}

///|
/// Connect, authenticate, and drive the socket until shutdown.
/// A second invocation fails even after the first has returned.
pub async fn ClientExecutor::run(self : ClientExecutor) -> Unit {
  if self.shared.started.val {
    raise ClientError::ExecutorAlreadyStarted
  }
  self.shared.started.val = true
  if self.shared.closed.val {
    self.shared.ready.complete(
      Some(ClientError::Closed("connection is closed")),
    )
    self.shared.done.complete(None)
    return
  }
  let mut failure : Error? = None
  defer {
    let terminal_error : Error? = if !self.shared.ready.done.val {
      if self.shared.aborted.val {
        Some(ClientError::Closed("connection aborted"))
      } else if failure is Some(error) {
        Some(error)
      } else {
        Some(ClientError::Closed("executor stopped before startup"))
      }
    } else {
      failure
    }
    let runtime_close_error : Error? = if self.shared.aborted.val {
      Some(ClientError::Closed("connection aborted"))
    } else {
      terminal_error
    }
    self.shared.driver_task.val = None
    self.shared.ready.complete(terminal_error)
    close_runtime(self.shared, error?=runtime_close_error)
    if self.shared.aborted.val {
      self.shared.control.close(error=ClientError::Closed("connection aborted"))
    } else {
      self.shared.control.close()
    }
    self.shared.done.complete(failure)
  }
  @async.with_task_group(group => {
    let task = group.spawn(no_wait=true, allow_failure=true, () => {
      self.run_connection(group)
    })
    self.shared.driver_task.val = Some(task)
    task.wait()
  }) catch {
    @async.WaitedTaskAlreadyCancelled as error =>
      if !self.shared.aborted.val && !self.shared.close_requested.val {
        failure = Some(error)
      }
    error => if !self.shared.aborted.val { failure = Some(error) }
  }
  if failure is Some(error) {
    raise error
  }
}

///|
async fn ClientExecutor::run_connection(
  self : ClientExecutor,
  group : @async.TaskGroup[Unit],
) -> Unit {
  if self.shared.closing.val {
    return
  }
  defer (if self.shared.aborted.val {
    self.shared.control.close(error=ClientError::Closed("connection aborted"))
  } else {
    self.shared.control.close()
  })
  group.spawn_bg(() => run_control_queue(self.shared.control, group))
  let config = self.shared.config
  let (setup, state) = if config.connect_timeout_ms is Some(timeout_ms) {
    @async.with_timeout(
      timeout_ms,
      () => connect_authenticated(config),
      error=ClientError::Protocol("connection timed out"),
    )
  } else {
    connect_authenticated(config)
  }
  let transferred = Ref(false)
  defer (if !transferred.val { setup.stream.close() })
  if self.shared.closing.val {
    return
  }
  for name, value in state.parameters {
    self.shared.parameters[name] = value
  }
  self.shared.process_id.val = state.process_id
  self.shared.secret_key.val = state.secret_key
  self.shared.transaction_status.val = state.transaction_status
  self.shared.ready.complete(None)
  transferred.val = true
  let driver : Driver = { shared: self.shared, stream: setup.stream, }
  driver.run(group)
  match self.shared.runtime_error.val {
    Some(error) if !self.shared.aborted.val => raise error
    _ => ()
  }
}

///|
async fn connect_authenticated(
  config : Config,
) -> (ConnectionSetup, StartupState) {
  let setup = connect_stream(config)
  let transferred = Ref(false)
  defer (if !transferred.val { setup.stream.close() })
  let state = startup(setup, config)
  transferred.val = true
  (setup, state)
}

///|
async fn run_control_queue(
  control : DrainControl,
  group : @async.TaskGroup[Unit],
) -> Unit {
  while true {
    let work = control.queue.get() catch {
      @aqueue.QueueAlreadyClosed => return
      error => raise error
    }
    group.spawn_bg(work)
  }
}

///|
/// Wait for startup. All waiters observe the same result.
pub async fn Client::ready(self : Client) -> Unit {
  self.shared.ready.wait()
}

///|
/// Return whether the shared runtime can no longer accept or perform work.
///
/// Graceful close sets this after the driver exits. Abort sets it immediately;
/// the driver then closes the physical transport while cancellation unwinds.
pub fn Client::is_closed(self : Client) -> Bool {
  self.shared.closed.val
}

///|
/// Return whether the latest `ReadyForQuery` reported an idle transaction.
///
/// This is a snapshot of the most recent protocol response, not a live server
/// check. It returns `false` before successful startup and after the client
/// closes, even though the initial stored status is `I`.
pub fn Client::is_transaction_idle(self : Client) -> Bool {
  self.shared.ready.done.val &&
  self.shared.ready.error.val is None &&
  !self.shared.closed.val &&
  self.shared.transaction_status.val == b'I'
}

///|
/// Gracefully shut down the client runtime.
///
/// Close is an operation barrier: it waits behind work already accepted by the
/// FIFO gate, rejects later work, sends `Terminate`, and returns only after the
/// private driver has closed the transport. Cancellation of the caller still
/// propagates while executor failure after shutdown is ignored.
pub async fn Client::close(self : Client) -> Unit {
  let permit = self.shared.operation_gate.acquire()
  defer permit.release()
  if self.shared.closing.val {
    self.wait_for_shutdown()
    return
  }
  self.shared.close_requested.val = true
  self.shared.closing.val = true
  if !self.shared.ready.done.val {
    if self.shared.driver_task.val is Some(task) {
      task.cancel()
    } else {
      self.shared.ready.complete(
        Some(ClientError::Closed("connection is closed")),
      )
      close_runtime(self.shared)
      self.shared.control.close()
      self.shared.done.complete(None)
    }
    self.wait_for_shutdown()
    return
  }
  let responses = @async.Queue(kind=Unbounded)
  let submitted = self.shared.requests.try_put({
    kind: Terminate,
    bytes: terminate_bytes(),
    responses,
    copy_input: None,
    copy_error: None,
  }) catch {
    @aqueue.QueueAlreadyClosed => {
      close_runtime(self.shared)
      false
    }
    error => {
      close_runtime(self.shared, error~)
      false
    }
  }
  if !submitted {
    responses.close()
    return
  }
  self.wait_for_shutdown()
}

///|
/// Wait for transport shutdown while keeping executor failure separate from
/// cancellation of the caller that is closing the client.
async fn Client::wait_for_shutdown(self : Client) -> Unit {
  self.shared.done.wait() catch {
    @async.WaitedTaskAlreadyCancelled as error => raise error
    _ => ()
  }
}

///|
/// Immediately abandon this physical connection.
///
/// Unlike `close`, abort does not wait behind active work or send PostgreSQL's
/// graceful `Terminate` message. It synchronously rejects new work and requests
/// driver cancellation; the driver closes the transport as cancellation
/// unwinds. Active and queued operations fail with
/// `ClientError::Closed("connection aborted")`. Repeated calls are safe.
pub fn Client::abort(self : Client) -> Unit {
  if self.shared.closed.val || self.shared.aborted.val {
    return
  }
  self.shared.aborted.val = true
  self.shared.closing.val = true
  self.shared.closed.val = true
  let error = ClientError::Closed("connection aborted")
  // Preserve buffered requests so the driver's teardown can close every
  // request-local response queue with the same abort error.
  self.shared.requests.close(error~)
  self.shared.async_messages.close()
  if self.shared.driver_task.val is Some(task) {
    task.cancel()
  } else {
    self.shared.ready.complete(Some(error))
    self.shared.control.close(error~)
    self.shared.done.complete(None)
  }
}

///|
/// Create a token capable of cancelling the current backend operation.
///
/// The token stays valid until the underlying connection closes. It can be
/// stored and used from another task or timeout handler when a long-running
/// query should be interrupted.
pub fn Client::cancel_token(self : Client) -> CancelToken raise {
  if !self.shared.ready.done.val || self.shared.ready.error.val is Some(_) {
    raise ClientError::NotReady
  }
  {
    config: self.shared.config,
    process_id: self.shared.process_id,
    secret_key: self.shared.secret_key,
    closed: self.shared.closed,
  }
}

///|
/// Return the backend process ID embedded in the cancellation token.
pub fn CancelToken::process_id(self : CancelToken) -> Int32 {
  self.process_id.val
}

///|
/// Reset the shared type cache to the built-in PostgreSQL descriptors only.
///
/// Use this when the application expects server-side type definitions to change
/// and wants later queries to re-fetch catalog metadata lazily.
pub fn Client::clear_type_cache(self : Client) -> Unit {
  self.shared.types.clear()
  seed_builtin_types(self.shared.types)
}

///|
/// Look up a server parameter learned during startup or later async updates.
pub fn Client::parameter(self : Client, name : String) -> String? raise {
  if !self.shared.ready.done.val || self.shared.ready.error.val is Some(_) {
    raise ClientError::NotReady
  }
  self.shared.parameters.get(name)
}

///|
/// Perform a cheap round trip that validates the connection is still usable.
///
/// The method sends a bare `Sync`, waits for the matching `ReadyForQuery`, and
/// therefore confirms that both the outbound request path and inbound response
/// path are still alive.
pub async fn Client::check_connection(self : Client) -> Unit {
  self.run_operation(context => {
    let responses = context.client.send_request(Messages, sync_bytes())
    drain_sync_response(responses)
  })
}

///|
/// Send a PostgreSQL cancel request over a short-lived control connection.
///
/// PostgreSQL requires cancellation to happen on a separate TCP connection that
/// carries only the backend process ID and secret key from startup.
pub async fn CancelToken::cancel(self : CancelToken) -> Unit {
  if self.closed.val {
    return
  }
  let conn = @socket.Tcp::connect_to_host(
    self.config.connect_host(),
    port=self.config.port,
  )
  defer conn.close()
  let buf = Buffer()
  @frontend.cancel_request(self.process_id.val, self.secret_key.val, buf)
  conn.write(buf.to_bytes())
}