///|
/// An incremental tool-call fragment within a streaming chunk.
///
/// During streaming, a tool call is delivered piecewise: the `index`
/// identifies which call it belongs to, and any of `id` / `name` /
/// `arguments` may be a partial fragment to be concatenated.
pub(all) struct ToolCallDelta {
  index : Int
  id : String?
  name : String?
  arguments : String?
} derive(Eq, Debug)

///|
/// A parsed delta from a streaming chat completion chunk.
pub(all) struct StreamChunk {
  /// Incremental text content, if any.
  content : String?
  /// Incremental tool-call fragments, if any.
  tool_calls : Array[ToolCallDelta]
  /// The finish reason, present on the final chunk of a choice.
  finish_reason : String?
} derive(Eq, Debug)

///|
/// Parse one SSE `data:` payload (the JSON after `data: `) into a `StreamChunk`.
///
/// Returns `None` for payloads that carry no usable delta (e.g. role-only
/// opening chunks). Raises `LLMError::Decode` on malformed JSON.
pub fn parse_stream_data(payload : String) -> StreamChunk? raise LLMError {
  let json = @json.parse(payload) catch {
    err => raise Decode("stream chunk: " + err.to_string())
  }
  guard json is Object(obj) else { return None }
  guard obj.get("choices") is Some(Array(choices)) else { return None }
  guard choices.get(0) is Some(Object(choice)) else { return None }
  let delta = match choice.get("delta") {
    Some(Object(d)) => d
    _ => {}
  }
  let content = match delta.get("content") {
    Some(String(s)) => Some(s)
    _ => None
  }
  let tool_calls = parse_tool_call_deltas(delta)
  let finish_reason = match choice.get("finish_reason") {
    Some(String(s)) => Some(s)
    _ => None
  }
  if content is None && tool_calls.length() == 0 && finish_reason is None {
    return None
  }
  Some({ content, tool_calls, finish_reason })
}

///|
/// Extract tool-call fragments from a streaming `delta` object.
fn parse_tool_call_deltas(delta : Map[String, Json]) -> Array[ToolCallDelta] {
  let out = []
  guard delta.get("tool_calls") is Some(Array(arr)) else { return out }
  for item in arr {
    guard item is Object(tc) else { continue }
    let index = match tc.get("index") {
      Some(Number(n, ..)) => n.to_int()
      _ => 0
    }
    let id = match tc.get("id") {
      Some(String(s)) => Some(s)
      _ => None
    }
    let (name, arguments) = match tc.get("function") {
      Some(Object(f)) => {
        let name = match f.get("name") {
          Some(String(s)) => Some(s)
          _ => None
        }
        let args = match f.get("arguments") {
          Some(String(s)) => Some(s)
          _ => None
        }
        (name, args)
      }
      _ => (None, None)
    }
    out.push({ index, id, name, arguments })
  }
  out
}

///|
/// Accumulates streaming `ToolCallDelta` fragments into complete `ToolCall`s.
///
/// Feed each chunk's `tool_calls` via `add`; call `finish` to obtain the
/// assembled list once the stream completes.
pub struct ToolCallAccumulator {
  by_index : Map[Int, ToolCallBuilder]
  order : Array[Int]
}

///|
/// Internal mutable builder for one in-progress tool call.
struct ToolCallBuilder {
  mut id : String
  mut name : String
  arguments : StringBuilder
}

///|
/// Create an empty accumulator.
pub fn ToolCallAccumulator::new() -> ToolCallAccumulator {
  { by_index: {}, order: [] }
}

///|
/// Fold a chunk's tool-call fragments into the accumulator.
pub fn ToolCallAccumulator::add(
  self : ToolCallAccumulator,
  deltas : Array[ToolCallDelta],
) -> Unit {
  for d in deltas {
    let builder = match self.by_index.get(d.index) {
      Some(b) => b
      None => {
        let b = { id: "", name: "", arguments: StringBuilder::new() }
        self.by_index[d.index] = b
        self.order.push(d.index)
        b
      }
    }
    if d.id is Some(id) {
      builder.id = id
    }
    if d.name is Some(name) {
      builder.name = name
    }
    if d.arguments is Some(args) {
      builder.arguments.write_string(args)
    }
  }
}

///|
/// Assemble the accumulated fragments into complete tool calls, in the order
/// they first appeared.
pub fn ToolCallAccumulator::finish(
  self : ToolCallAccumulator,
) -> Array[ToolCall] {
  let out = []
  for idx in self.order {
    if self.by_index.get(idx) is Some(b) {
      out.push({
        id: b.id,
        function: { name: b.name, arguments: b.arguments.to_string() },
      })
    }
  }
  out
}

///|
/// Strip a single leading SSE field prefix (`data:` / `event:` ...) and the
/// optional single space after the colon, returning the field value.
///
/// Returns `None` for comment lines (starting with `:`) and blank lines.
pub fn parse_sse_line(line : String) -> (String, String)? {
  // Strip a trailing "\n" and/or "\r" left by the line framing.
  let mut end = line.length()
  if end > 0 && line[end - 1] == '\n' {
    end = end - 1
  }
  if end > 0 && line[end - 1] == '\r' {
    end = end - 1
  }
  let trimmed = line[:end].to_owned()
  if trimmed.length() == 0 || trimmed.has_prefix(":") {
    return None
  }
  match trimmed.find(":") {
    Some(idx) => {
      let field = trimmed[:idx].to_owned()
      let mut rest = trimmed[idx + 1:].to_owned()
      // SSE allows a single optional space after the colon.
      if rest.has_prefix(" ") {
        rest = rest[1:].to_owned()
      }
      Some((field, rest))
    }
    None => Some((trimmed, ""))
  }
}

///|
/// The complete result of a streaming chat completion.
pub(all) struct StreamResult {
  /// The full accumulated assistant text.
  content : String
  /// Any tool calls assembled from streamed fragments.
  tool_calls : Array[ToolCall]
  /// The finish reason of the last chunk, if seen.
  finish_reason : String?
}

///|
/// Perform a streaming chat completion, invoking `on_chunk` for every parsed
/// `StreamChunk` as it arrives, and returning the fully assembled result
/// (text + tool calls) once the stream completes.
///
/// This is the low-level streaming entry point. For the common
/// text-only case, prefer `chat_stream`.
///
/// Raises `LLMError` on transport failure, non-2xx status, or decode failure.
pub async fn Client::chat_stream_full(
  self : Client,
  request : ChatRequest,
  on_chunk : (StreamChunk) -> Unit,
) -> StreamResult raise LLMError {
  request.stream = true
  let url = self.endpoint("/chat/completions")
  let headers = {
    "Authorization": "Bearer " + self.api_key,
    "Content-Type": "application/json",
    "Accept": "text/event-stream",
  }
  let body = request.to_json().stringify()
  let client = @mio.post_stream(url, headers~) catch {
    err => raise Transport(err.to_string())
  }
  // Ensure the underlying connection is always released.
  defer client.close()
  let content = StringBuilder::new()
  let tools = ToolCallAccumulator::new()
  let mut finish_reason : String? = None
  let parser = SSEParser::new()
  let response = try {
    client.write(body)
    client.end_request()
  } catch {
    err => raise Transport(err.to_string())
  }
  guard response.code >= 200 && response.code < 300 else {
    let rest = client.read_all().binary().to_unchecked_string() catch {
        _ => ""
      }
    raise ApiError(code=response.code, message=rest)
  }
  // Read the body line by line, framing SSE events with the parser.
  let mut done = false
  for ;; {
    if done {
      break
    }
    let line = client.read_until("\n") catch {
      err => raise Stream(err.to_string())
    }
    match line {
      None => break // EOF
      Some(l) =>
        if parser.push_line(l) is Some(event) {
          let payload = event.data
          if payload == "[DONE]" {
            done = true
          } else if parse_stream_data(payload) is Some(chunk) {
            if chunk.content is Some(text) {
              content.write_string(text)
            }
            if chunk.tool_calls.length() > 0 {
              tools.add(chunk.tool_calls)
            }
            if chunk.finish_reason is Some(_) {
              finish_reason = chunk.finish_reason
            }
            on_chunk(chunk)
          }
        }
    }
  }
  { content: content.to_string(), tool_calls: tools.finish(), finish_reason }
}

///|
/// Perform a streaming chat completion (text only).
///
/// `on_delta` is invoked for each incremental text fragment as it arrives.
/// Returns the fully accumulated assistant text once the stream completes.
///
/// Raises `LLMError` on transport failure, non-2xx status, or decode failure.
pub async fn Client::chat_stream(
  self : Client,
  request : ChatRequest,
  on_delta : (String) -> Unit,
) -> String raise LLMError {
  let result = self.chat_stream_full(request, fn(chunk) {
    if chunk.content is Some(text) {
      on_delta(text)
    }
  })
  result.content
}