///|
pub struct Function {
  arguments : String
  name : String
}

///|
impl ToJson for Function with to_json(self : Function) -> Json {
  { "name": self.name, "arguments": self.arguments }
}

///|
impl @json.FromJson for Function with from_json(
  json : Json,
  json_path : @json.JsonPath,
) -> Function {
  guard json is Object(json) else {
    raise @json.JsonDecodeError((json_path, "Expected an object"))
  }
  guard json.get("arguments") is Some(String(arguments)) else {
    raise @json.JsonDecodeError((json_path, "Missing 'arguments' field"))
  }
  guard json.get("name") is Some(String(name)) else {
    raise @json.JsonDecodeError((json_path, "Missing 'name' field"))
  }
  Function::{ arguments, name }
}

///|
pub struct ChatCompletionMessageToolCall {
  id : String
  function : Function
}

///|
impl ToJson for ChatCompletionMessageToolCall with to_json(
  self : ChatCompletionMessageToolCall,
) -> Json {
  { "id": self.id, "function": self.function, "type": "function" }
}

///|
impl @json.FromJson for ChatCompletionMessageToolCall with from_json(
  json : Json,
  json_path : @json.JsonPath,
) -> ChatCompletionMessageToolCall {
  guard json is Object(json) else {
    raise @json.JsonDecodeError((json_path, "Expected an object"))
  }
  guard json.get("id") is Some(String(id)) else {
    raise @json.JsonDecodeError((json_path, "Missing 'id' field"))
  }
  guard json.get("function") is Some(function) else {
    raise @json.JsonDecodeError((json_path, "Missing 'function' field"))
  }
  let function : Function = @json.from_json(
    function,
    path=json_path.add_key("function"),
  )
  ChatCompletionMessageToolCall::{ id, function }
}

///|
pub struct ChatCompletionMessage {
  content : String?
  tool_calls : Array[ChatCompletionMessageToolCall]
}

///|
pub fn ChatCompletionMessage::to_param(
  self : ChatCompletionMessage,
) -> ChatCompletionMessageParam {
  Assistant({ content: self.content, name: None, tool_calls: self.tool_calls })
}

///|
impl @json.FromJson for ChatCompletionMessage with from_json(
  json : Json,
  json_path : @json.JsonPath,
) -> ChatCompletionMessage {
  guard json is Object(json) else {
    raise @json.JsonDecodeError((json_path, "Expected an object"))
  }
  let content = match json.get("content") {
    Some(content) => {
      let content : String = @json.from_json(
        content,
        path=json_path.add_key("content"),
      )
      Some(content)
    }
    None => None
  }
  let tool_calls = match json.get("tool_calls") {
    Some(tool_calls) => {
      let tool_calls : Array[ChatCompletionMessageToolCall] = @json.from_json(
        tool_calls,
        path=json_path.add_key("tool_calls"),
      )
      tool_calls
    }
    None => []
  }
  ChatCompletionMessage::{ content, tool_calls }
}

///|
pub enum ChatCompletionChoiceFinishReason {
  Stop
  Length
  ToolCalls
  ContentFilter
  FunctionCall
}

///|
impl @json.FromJson for ChatCompletionChoiceFinishReason with from_json(
  json : Json,
  json_path : @json.JsonPath,
) -> ChatCompletionChoiceFinishReason {
  guard json is String(reason) else {
    raise @json.JsonDecodeError((json_path, "Expected a string"))
  }
  match reason {
    "stop" => ChatCompletionChoiceFinishReason::Stop
    "length" => ChatCompletionChoiceFinishReason::Length
    "tool_calls" => ChatCompletionChoiceFinishReason::ToolCalls
    "content_filter" => ChatCompletionChoiceFinishReason::ContentFilter
    "function_call" => ChatCompletionChoiceFinishReason::FunctionCall
    _ =>
      raise @json.JsonDecodeError(
        (json_path, "Unknown finish reason: \{reason}"),
      )
  }
}

///|
pub struct ChatCompletionChoice {
  finish_reason : ChatCompletionChoiceFinishReason
  index : Int
  message : ChatCompletionMessage
}

///|
impl @json.FromJson for ChatCompletionChoice with from_json(
  json : Json,
  json_path : @json.JsonPath,
) -> ChatCompletionChoice {
  guard json is Object(json) else {
    raise @json.JsonDecodeError((json_path, "Expected an object"))
  }
  guard json.get("finish_reason") is Some(finish_reason) else {
    raise @json.JsonDecodeError((json_path, "Missing 'finish_reason' field"))
  }
  let finish_reason : ChatCompletionChoiceFinishReason = @json.from_json(
    finish_reason,
    path=json_path.add_key("finish_reason"),
  )
  guard json.get("index") is Some(Number(index, ..)) else {
    raise @json.JsonDecodeError((json_path, "Missing 'index' field"))
  }
  guard json.get("message") is Some(message) else {
    raise @json.JsonDecodeError((json_path, "Missing 'message' field"))
  }
  let message : ChatCompletionMessage = @json.from_json(
    message,
    path=json_path.add_key("message"),
  )
  ChatCompletionChoice::{ finish_reason, index: index.to_int(), message }
}

///|
pub struct CompletionUsage {
  completion_tokens : Int
  prompt_tokens : Int
  total_tokens : Int
}

///|
impl @json.FromJson for CompletionUsage with from_json(
  json : Json,
  json_path : @json.JsonPath,
) -> CompletionUsage {
  guard json is Object(json) else {
    raise @json.JsonDecodeError((json_path, "Expected an object"))
  }
  guard json.get("completion_tokens") is Some(Number(completion_tokens, ..)) else {
    raise @json.JsonDecodeError(
      (json_path, "Missing 'completion_tokens' field"),
    )
  }
  guard json.get("prompt_tokens") is Some(Number(prompt_tokens, ..)) else {
    raise @json.JsonDecodeError((json_path, "Missing 'prompt_tokens' field"))
  }
  guard json.get("total_tokens") is Some(Number(total_tokens, ..)) else {
    raise @json.JsonDecodeError((json_path, "Missing 'total_tokens' field"))
  }
  CompletionUsage::{
    completion_tokens: completion_tokens.to_int(),
    prompt_tokens: prompt_tokens.to_int(),
    total_tokens: total_tokens.to_int(),
  }
}

///|
pub struct ChatCompletion {
  id : String
  choices : Array[ChatCompletionChoice]
  created : Int
  model : String
  usage : CompletionUsage?
}

///|
impl @json.FromJson for ChatCompletion with from_json(
  json : Json,
  json_path : @json.JsonPath,
) -> ChatCompletion {
  guard json is Object(json) else {
    raise @json.JsonDecodeError((json_path, "Expected an object"))
  }
  guard json.get("id") is Some(String(id)) else {
    raise @json.JsonDecodeError((json_path, "Missing 'id' field"))
  }
  guard json.get("choices") is Some(choices) else {
    raise @json.JsonDecodeError((json_path, "Missing 'choices' field"))
  }
  let choices : Array[ChatCompletionChoice] = @json.from_json(
    choices,
    path=json_path.add_key("choices"),
  )
  guard json.get("created") is Some(Number(created, ..)) else {
    raise @json.JsonDecodeError((json_path, "Missing 'created' field"))
  }
  guard json.get("model") is Some(String(model)) else {
    raise @json.JsonDecodeError((json_path, "Missing 'model' field"))
  }
  let usage = match json.get("usage") {
    Some(usage) => {
      let usage : CompletionUsage = @json.from_json(
        usage,
        path=json_path.add_key("usage"),
      )
      Some(usage)
    }
    None => None
  }
  ChatCompletion::{ id, choices, created: created.to_int(), model, usage }
}

///|
pub struct ChatCompletionSystemMessageParam {
  content : String
  name : String?
}

///|
pub struct ChatCompletionUserMessageParam {
  content : String
  name : String?
}

///|
typealias ChatCompletionMessageToolCall as ChatCompletionMessageToolCallParam

///|
pub struct ChatCompletionAssistantMessageParam {
  content : String?
  name : String?
  tool_calls : Array[ChatCompletionMessageToolCallParam]
}

///|
pub struct ChatCompletionToolMessageParam {
  content : String
  tool_call_id : String
}

///|
pub enum ChatCompletionMessageParam {
  System(ChatCompletionSystemMessageParam)
  User(ChatCompletionUserMessageParam)
  Assistant(ChatCompletionAssistantMessageParam)
  Tool(ChatCompletionToolMessageParam)
}

///|
impl ToJson for ChatCompletionMessageParam with to_json(
  self : ChatCompletionMessageParam,
) {
  match self {
    System(param) => {
      let json : Map[String, Json] = {
        "role": "system",
        "content": Json::string(param.content),
      }
      if param.name is Some(name) {
        json["name"] = Json::string(name)
      }
      Json::object(json)
    }
    User(param) => {
      let json : Map[String, Json] = {
        "role": "user",
        "content": Json::string(param.content),
      }
      if param.name is Some(name) {
        json["name"] = Json::string(name)
      }
      Json::object(json)
    }
    Assistant(param) => {
      let json : Map[String, Json] = {
        "role": "assistant",
        "tool_calls": param.tool_calls.to_json(),
      }
      if param.content is Some(content) {
        json["content"] = Json::string(content)
      }
      if param.name is Some(name) {
        json["name"] = Json::string(name)
      }
      Json::object(json)
    }
    Tool(param) =>
      {
        "role": "tool",
        "content": param.content,
        "tool_call_id": param.tool_call_id,
      }
  }
}

///|
pub fn system_message(
  content~ : String,
  name? : String,
) -> ChatCompletionMessageParam {
  ChatCompletionMessageParam::System({ content, name })
}

///|
pub fn user_message(
  content~ : String,
  name? : String,
) -> ChatCompletionMessageParam {
  ChatCompletionMessageParam::User({ content, name })
}

///|
pub fn tool_message(
  content~ : String,
  tool_call_id~ : String,
) -> ChatCompletionMessageParam {
  ChatCompletionMessageParam::Tool({ content, tool_call_id })
}

///|
pub fn Client::create_chat_completion(
  self : Client,
  messages~ : Array[ChatCompletionMessageParam],
  model~ : String,
  tools~ : Array[ChatCompletionToolParam] = [],
) -> ChatCompletion raise {
  let messages = messages.to_json()
  let json : Json = {
    "model": model,
    "messages": messages,
    "tools": tools.to_json(),
  }
  let text = @encoding.encode(json.stringify(), encoding=UTF8)
  let request = @buffer.new()
  request.write_bytes("POST ")
  for segment in self.base_url.path {
    request.write_byte('/')
    request.write_bytesview(segment)
  }
  request.write_bytes("/chat/completions HTTP/1.1\r\n")
  request.write_bytes("Host: openrouter.ai\r\n")
  request.write_bytes("User-Agent: curl/8.7.1\r\n")
  request.write_bytes("Accept: */*\r\n")
  request.write_bytes("Content-Type: application/json\r\n")
  if self.api_key is Some(api_key) {
    request.write_bytes("Authorization: Bearer ")
    request.write_bytes(api_key)
    request.write_bytes("\r\n")
  }
  request.write_bytes("Connection: close\r\n")
  request.write_bytes("Content-Length: ")
  request.write_bytes(
    @encoding.encode(text.length().to_string(), encoding=UTF8),
  )
  request.write_bytes("\r\n")
  request.write_bytes("\r\n")
  request.write_bytes(text)
  let http = request.contents()
  let response = self.send_http_request(http)[:]
  let headers = []
  let mut header = response
  let body = loop response {
    [] as body => body
    [.. "\r\n\r\n", .. body] => body
    [.. "\r\n", .. k] as response => {
      headers.push(header[:response.start_offset() - header.start_offset()])
      header = k
      continue k
    }
    [_, .. k] => continue k
  }
  let mut chunked_encoding = false
  for header in headers {
    if header == "Transfer-Encoding: chunked" {
      chunked_encoding = true
    }
  }
  let body = if chunked_encoding {
    let buffer = @buffer.new()
    let mut chunk_size = 0
    loop body {
      [] => ()
      ['0'..='9' as c, .. body] => {
        chunk_size = chunk_size * 16 + (c.to_int() - '0')
        continue body
      }
      ['A'..='F' as c, .. body] => {
        chunk_size = chunk_size * 16 + (c.to_int() - 'A' + 10)
        continue body
      }
      ['a'..='f' as c, .. body] => {
        chunk_size = chunk_size * 16 + (c.to_int() - 'a' + 10)
        continue body
      }
      [.. "\r\n", .. body] => {
        let chunk_data = body[:chunk_size]
        buffer.write_bytesview(chunk_data)
        let body = body[chunk_size:]
        let body = match body {
          [.. "\r\n", .. body] => body
          [..] => fail("Invalid chunk data")
        }
        chunk_size = 0
        continue body
      }
      [..] as body =>
        fail(
          "Invalid chunk size: \{body |> @encoding.decode_lossy(_, encoding=UTF8)}",
        )
    }
    loop buffer.contents()[:] {
      [' ' | '\n' | '\r', .. k] => continue k
      [.. k] => k
    }
  } else {
    body
  }
  body
  |> @encoding.decode_lossy(encoding=UTF8)
  |> @json.parse()
  |> @json.from_json()
}

///|
priv struct FunctionDefinition {
  name : String
  description : String
  parameters : Map[String, Json]
  strict : Bool?
}

///|
impl ToJson for FunctionDefinition with to_json(self : FunctionDefinition) -> Json {
  let json : Map[String, Json] = {
    "name": Json::string(self.name),
    "description": Json::string(self.description),
    "parameters": Json::object(self.parameters),
  }
  if self.strict is Some(bool) {
    json["strict"] = Json::boolean(bool)
  }
  Json::object(json)
}

///|
pub fn function_tool(
  name~ : String,
  description~ : String,
  parameters~ : Map[String, Json],
  strict? : Bool,
) -> ChatCompletionToolParam {
  Function(FunctionDefinition::{ name, description, parameters, strict })
}

///|
enum ChatCompletionToolParam {
  Function(FunctionDefinition)
}

///|
impl ToJson for ChatCompletionToolParam with to_json(
  self : ChatCompletionToolParam,
) -> Json {
  match self {
    Function(function_definition) =>
      { "type": "function", "function": function_definition.to_json() }
  }
}