// The socket-owning connection task and its request/response state machine.

///|
/// Drive the PostgreSQL connection until graceful shutdown or failure.
///
/// The loop alternates between:
/// 1. taking queued client requests,
/// 2. opportunistically pipelining additional non-barrier requests, and
/// 3. forwarding backend messages either to async listeners or to the oldest
///    pending request.
///
/// On exit, including cancellation, the loop closes all outstanding queues and
/// the socket, and marks the shared runtime as closed. Ordinary errors are sent
/// to waiting requests; cancellation closes them with `ClientError::Closed`
/// while the driver task itself remains cancelled.
pub async fn Connection::run(
  self : Connection,
  on_async? : (AsyncMessage) -> Unit = _ => (),
) -> Unit {
  let pending : Array[Request] = []
  let mut cleaned_up = false
  fn cleanup(error : Error?) {
    if cleaned_up {
      return
    }
    cleaned_up = true
    let error = error.unwrap_or(ClientError::Closed("connection is closed"))
    self.close_pending_requests(pending, Some(error))
    self.drain_queued_requests(Some(error))
    close_runtime(self.shared, error~)
    self.stream.close()
  }
  // Also cover exit before the group's main task starts or while joining it.
  defer cleanup(None)
  @async.with_task_group(group => {
    self.shared.background_group.val = Some(group)
    defer {
      self.shared.background_group.val = None
      // Wake even cancellation-protected drains before the group joins them.
      cleanup(None)
    }

    let mut sent_terminate = false
    try {
      while true {
        if pending.is_empty() && !sent_terminate {
          // When nothing is in flight, block until the runtime has a request
          // that can own subsequent backend replies.
          let request = self.shared.requests.get()
          sent_terminate = self.send_outbound_request(request, pending)
        }
        if self.can_pipeline_requests(pending) && !sent_terminate {
          // If the current in-flight tail contains no barriers, batch more work
          // onto the socket before we wait for the next backend frame.
          sent_terminate = self.pipeline_pending_requests(pending)
        }
        if pending.is_empty() {
          if sent_terminate {
            // `Terminate` has no normal response stream. Once it is the only
            // outstanding work, shutdown can converge locally.
            return
          }
          continue
        }
        let message = self.read_message()
        // After async side-channel frames are peeled off, protocol order says
        // the oldest pending request owns the next reply.
        self.forward_pending_message(pending, message, on_async)
      }
    } catch {
      err => cleanup(Some(err))
    }
  })
}

///|
/// Receive the next asynchronous message buffered by the connection loop.
///
/// Returns `None` after the async-message queue is closed. This is usually read
/// from one dedicated task that owns async notices / notifications for the
/// connection.
pub async fn Connection::next_message(self : Connection) -> AsyncMessage? {
  Some(self.shared.async_messages.get()) catch {
    _ => None
  }
}

///|
/// Look up the latest value of a server parameter.
pub fn Connection::parameter(self : Connection, name : String) -> String? {
  self.shared.parameters.get(name)
}

///|
/// Write one queued request to the wire and decide whether it stops pipelining.
async fn Connection::send_outbound_request(
  self : Connection,
  request : Request,
  pending : Array[Request],
) -> Bool {
  // Once dequeued, a request must remain reachable by exit cleanup even if
  // writing its bytes fails or is cancelled before the write completes.
  pending.push(request)
  self.stream.write(request.bytes)
  match request.kind {
    Terminate => {
      // PostgreSQL closes the socket after `Terminate`; there is no matching
      // terminal `ReadyForQuery` for this synthetic request queue.
      ignore(pending.pop())
      request.responses.close()
      true
    }
    Messages | CopyIn => false
  }
}

///|
/// Return whether every currently pending request may be safely pipelined.
///
/// COPY IN and termination requests act as barriers because they require
/// request-specific coordination with subsequent protocol messages.
fn Connection::can_pipeline_requests(
  self : Connection,
  pending : Array[Request],
) -> Bool {
  ignore(self)
  for request in pending {
    match request.kind {
      Messages => ()
      _ => return false
    }
  }
  true
}

///|
/// Opportunistically enqueue more non-barrier requests without blocking.
async fn Connection::pipeline_pending_requests(
  self : Connection,
  pending : Array[Request],
) -> Bool {
  while true {
    let request = self.shared.requests.try_get() catch { _ => None }
    match request {
      None => return false
      Some(request) => {
        let sent_terminate = self.send_outbound_request(request, pending)
        // Stop as soon as pipelining reaches a request that needs extra
        // mid-stream coordination, such as COPY IN or graceful shutdown.
        let is_barrier = match request.kind {
          Messages => false
          _ => true
        }
        if sent_terminate || is_barrier {
          return sent_terminate
        }
        continue
      }
    }
  }
  false
}

///|
/// Route one backend message to either async listeners or the oldest request.
async fn Connection::forward_pending_message(
  self : Connection,
  pending : Array[Request],
  message : @backend.Message,
  on_async : (AsyncMessage) -> Unit,
) -> Unit {
  if self.handle_async_message(message, on_async) {
    // Notices, notifications, and parameter updates never advance a request.
    return
  }
  guard !pending.is_empty() else {
    raise ClientError::UnexpectedMessage("unexpected backend message")
  }
  let request = pending[0]
  // All remaining frames belong to the request at the front of the protocol
  // timeline because PostgreSQL replies in request order.
  match request.kind {
    Messages => {
      self.forward_message_to_request(request, message)
      if message is ReadyForQuery(_) {
        ignore(pending.remove(0))
      }
    }
    CopyIn => {
      self.forward_copy_in_message(request, message)
      if message is ReadyForQuery(_) {
        ignore(pending.remove(0))
      }
    }
    Terminate =>
      raise ClientError::UnexpectedMessage("unexpected terminate request")
  }
}

///|
/// Forward a regular request message to its response queue.
///
/// `ReadyForQuery` also updates the shared transaction status snapshot.
async fn Connection::forward_message_to_request(
  self : Connection,
  request : Request,
  message : @backend.Message,
) -> Unit {
  match message {
    ReadyForQuery(body) => {
      self.shared.transaction_status.val = body.status
      request.responses.put(message)
    }
    _ => request.responses.put(message)
  }
}

///|
/// Forward a COPY IN request message and start draining client input when ready.
async fn Connection::forward_copy_in_message(
  self : Connection,
  request : Request,
  message : @backend.Message,
) -> Unit {
  let copy_input = request.copy_input.unwrap()
  match message {
    CopyInResponse(_) => {
      request.responses.put(message)
      // Only after PostgreSQL enters COPY mode may the client start producing
      // `CopyData` / `CopyDone` / `CopyFail` frames for this request.
      self.drain_copy_input(copy_input)
    }
    ReadyForQuery(body) => {
      self.shared.transaction_status.val = body.status
      request.responses.put(message)
    }
    _ => request.responses.put(message)
  }
}

///|
/// Drain one COPY IN sink until it finishes or aborts.
///
/// The driver writes each client-produced action as a protocol message on the
/// shared connection stream. Once `Finish` or `Fail` is seen, the COPY request
/// returns to normal backend-response processing.
async fn Connection::drain_copy_input(
  self : Connection,
  input : @async.Queue[CopyInAction],
) -> Unit {
  for act = input.get() {
    match act {
      Data(data) => {
        // Data chunks can flow indefinitely; keep consuming producer input
        // until it explicitly finishes or aborts this COPY exchange.
        let payload = Buffer()
        @frontend.CopyData::new(data[:]).write(payload)
        self.stream.write(payload.to_bytes())
        continue input.get()
      }
      Finish => {
        // After `CopyDone`, the request returns to normal backend completion
        // handling on the response queue.
        let payload = Buffer()
        @frontend.copy_done(payload)
        self.stream.write(payload.to_bytes())
        return
      }
      Fail(message) => {
        // PostgreSQL still sends trailing completion frames after `CopyFail`,
        // so the request stays pending until the backend catches up.
        let payload = Buffer()
        @frontend.copy_fail(@proto.utf8_encode(message)[:], payload)
        self.stream.write(payload.to_bytes())
        return
      }
    }
  }
}

///|
/// Close one request's queues, optionally propagating an error.
fn Connection::close_request(
  self : Connection,
  request : Request,
  error : Error?,
) -> Unit {
  ignore(self)
  match error {
    None => request.responses.close()
    Some(err) => request.responses.close(error=err, clear=true)
  }
  match request.copy_input {
    Some(input) =>
      match error {
        None => input.close()
        Some(err) => input.close(error=err, clear=true)
      }
    None => ()
  }
}

///|
/// Close every currently pending request.
fn Connection::close_pending_requests(
  self : Connection,
  pending : Array[Request],
  error : Error?,
) -> Unit {
  for request in pending {
    self.close_request(request, error)
  }
}

///|
/// Drain and close requests that were queued but never written to the wire.
fn Connection::drain_queued_requests(self : Connection, error : Error?) -> Unit {
  let close_error = match error {
    Some(err) => Some(err)
    None => Some(ClientError::Closed("connection is closed"))
  }
  // Requests still sitting in the submission queue were never written to the
  // socket, so they can only fail locally as shutdown fallout.
  while true {
    let request = self.shared.requests.try_get() catch { _ => None }
    match request {
      None => return
      Some(request) => {
        self.close_request(request, close_error)
        continue
      }
    }
  }
}

///|
/// Handle backend messages that are not tied to the current request head.
///
/// PostgreSQL may send notices, notifications, and parameter updates while
/// ordinary requests are in flight. The connection loop peels those off first
/// so request-specific stream parsers only see messages that belong to them.
fn Connection::handle_async_message(
  self : Connection,
  message : @backend.Message,
  on_async : (AsyncMessage) -> Unit,
) -> Bool raise {
  match message {
    ParameterStatus(body) => {
      let name = body.name_str()
      let value = body.value_str()
      self.shared.parameters[name] = value
      self.push_async_message(ParameterStatus(name, value), on_async)
      true
    }
    NoticeResponse(body) => {
      self.push_async_message(
        Notice(parse_database_error(body.fields())),
        on_async,
      )
      true
    }
    NotificationResponse(body) => {
      self.push_async_message(
        Notification({
          process_id: body.process_id,
          channel: body.channel_str(),
          payload: body.message_str(),
        }),
        on_async,
      )
      true
    }
    _ => false
  }
}

///|
/// Publish one async message to both the callback and the queue buffer.
fn Connection::push_async_message(
  self : Connection,
  message : AsyncMessage,
  on_async : (AsyncMessage) -> Unit,
) -> Unit {
  on_async(message)
  ignore(self.shared.async_messages.try_put(message)) catch {
    _ => ()
  }
}

///|
/// Read and parse one backend message from the shared connection stream.
async fn Connection::read_message(self : Connection) -> @backend.Message {
  let header = self.stream.read_exactly(5)
  let body_len = @proto.ByteReader::new(header[1:]).read_i32_be()
  let body = self.stream.read_exactly(body_len - 4)
  let packet = Buffer(size_hint=header.length() + body.length())
  packet.write_bytes(header)
  packet.write_bytes(body)
  match @backend.Message::parse(packet.to_bytes()[:]) {
    Some(parsed) => parsed.message
    None => raise ClientError::UnexpectedMessage("incomplete backend message")
  }
}