///|
/// The role of a chat message author.
pub(all) enum Role {
  System
  User
  Assistant
  Tool
} derive(Eq, Debug)

///|
pub fn Role::to_string(self : Role) -> String {
  match self {
    System => "system"
    User => "user"
    Assistant => "assistant"
    Tool => "tool"
  }
}

///|
pub fn Role::parse(s : String) -> Role {
  match s {
    "system" => System
    "assistant" => Assistant
    "tool" => Tool
    _ => User
  }
}

///|
/// A single part of a multimodal message content.
///
/// OpenAI-compatible: a message's `content` may be a plain string or an array
/// of content parts. `Text` and `ImageUrl` cover the common cases.
pub(all) enum ContentPart {
  Text(String)
  ImageUrl(String)
} derive(Eq, Debug)

///|
/// The content of a message: either a single text string or a list of parts.
pub(all) enum Content {
  Str(String)
  Parts(Array[ContentPart])
} derive(Eq, Debug)

///|
/// The plain-text view of the content: the string itself, or the text parts
/// of a multimodal message concatenated together.
pub fn Content::to_text(self : Content) -> String {
  match self {
    Str(s) => s
    Parts(parts) => {
      let buf = StringBuilder::new()
      for p in parts {
        if p is Text(t) {
          buf.write_string(t)
        }
      }
      buf.to_string()
    }
  }
}

///|
pub impl Show for Content with fn output(self : Content, logger : &Logger) -> Unit {
  logger.write_string(self.to_text())
}

///|
pub impl ToJson for Content with fn to_json(self : Content) -> Json {
  match self {
    Str(s) => Json::string(s)
    Parts(parts) => {
      let arr = []
      for p in parts {
        match p {
          Text(t) =>
            arr.push(
              Json::object({
                "type": Json::string("text"),
                "text": Json::string(t),
              }),
            )
          ImageUrl(u) =>
            arr.push(
              Json::object({
                "type": Json::string("image_url"),
                "image_url": Json::object({ "url": Json::string(u) }),
              }),
            )
        }
      }
      Json::array(arr)
    }
  }
}

///|
pub impl @json.FromJson for Content with fn from_json(
  json : Json,
  path : @json.JsonPath,
) -> Content {
  match json {
    String(s) => Str(s)
    Array(arr) => {
      let parts = []
      for item in arr {
        match item {
          Object(obj) =>
            match obj.get("type") {
              Some(String("text")) =>
                match obj.get("text") {
                  Some(String(t)) => parts.push(Text(t))
                  _ => ()
                }
              Some(String("image_url")) =>
                match obj.get("image_url") {
                  Some(Object(iu)) =>
                    match iu.get("url") {
                      Some(String(u)) => parts.push(ImageUrl(u))
                      _ => ()
                    }
                  _ => ()
                }
              _ => ()
            }
          _ => ()
        }
      }
      Parts(parts)
    }
    _ =>
      raise @json.JsonDecodeError((path, "Content: expected string or array"))
  }
}

///|
/// A chat message.
pub(all) struct Message {
  role : Role
  content : Content
  /// Present on assistant messages that request tool calls.
  tool_calls : Array[ToolCall]?
  /// Present on `tool` role messages, linking the result to a prior call.
  tool_call_id : String?
  /// Optional author name.
  name : String?
} derive(Eq, Debug)

///|
pub impl ToJson for Message with fn to_json(self : Message) -> Json {
  let obj : Map[String, Json] = {
    "role": Json::string(self.role.to_string()),
    "content": self.content.to_json(),
  }
  if self.tool_calls is Some(tcs) {
    obj["tool_calls"] = tcs.to_json()
  }
  if self.tool_call_id is Some(id) {
    obj["tool_call_id"] = Json::string(id)
  }
  if self.name is Some(n) {
    obj["name"] = Json::string(n)
  }
  Json::object(obj)
}

///|
pub impl @json.FromJson for Message with fn from_json(
  json : Json,
  path : @json.JsonPath,
) -> Message {
  guard json is Object(obj) else {
    raise @json.JsonDecodeError((path, "Message: expected object"))
  }
  let role = match obj.get("role") {
    Some(String(r)) => Role::parse(r)
    _ => User
  }
  let content = match obj.get("content") {
    Some(c) => @json.from_json(c)
    None => Str("")
  }
  let tool_calls = match obj.get("tool_calls") {
    Some(Array(_) as tc) => Some(@json.from_json(tc))
    _ => None
  }
  let tool_call_id = match obj.get("tool_call_id") {
    Some(String(s)) => Some(s)
    _ => None
  }
  let name = match obj.get("name") {
    Some(String(s)) => Some(s)
    _ => None
  }
  { role, content, tool_calls, tool_call_id, name }
}

///|
/// Token usage statistics returned by the API.
pub(all) struct Usage {
  prompt_tokens : Int
  completion_tokens : Int
  total_tokens : Int
} derive(Eq, Debug, ToJson, FromJson)

///|
/// Add two usage records together (useful when aggregating across calls).
pub fn Usage::add(self : Usage, other : Usage) -> Usage {
  {
    prompt_tokens: self.prompt_tokens + other.prompt_tokens,
    completion_tokens: self.completion_tokens + other.completion_tokens,
    total_tokens: self.total_tokens + other.total_tokens,
  }
}

///|
/// A zero-valued usage record.
pub fn Usage::zero() -> Usage {
  { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }
}

///|
/// A function invocation requested by the model.
pub(all) struct FunctionCall {
  name : String
  /// Raw JSON-encoded argument string, exactly as returned by the API.
  arguments : String
} derive(Eq, Debug, ToJson, FromJson)

///|
/// A tool call requested by the assistant.
pub(all) struct ToolCall {
  id : String
  function : FunctionCall
} derive(Eq, Debug)

///|
pub impl ToJson for ToolCall with fn to_json(self : ToolCall) -> Json {
  Json::object({
    "id": Json::string(self.id),
    "type": Json::string("function"),
    "function": self.function.to_json(),
  })
}

///|
pub impl @json.FromJson for ToolCall with fn from_json(
  json : Json,
  path : @json.JsonPath,
) -> ToolCall {
  guard json is Object(obj) else {
    raise @json.JsonDecodeError((path, "ToolCall: expected object"))
  }
  let id = match obj.get("id") {
    Some(String(s)) => s
    _ => ""
  }
  let function = match obj.get("function") {
    Some(f) => @json.from_json(f)
    None => raise @json.JsonDecodeError((path, "ToolCall: missing function"))
  }
  { id, function }
}

///|
/// A single choice in a chat completion response.
pub(all) struct Choice {
  index : Int
  message : Message
  finish_reason : String?
} derive(Eq, Debug)

///|
pub impl @json.FromJson for Choice with fn from_json(
  json : Json,
  path : @json.JsonPath,
) -> Choice {
  guard json is Object(obj) else {
    raise @json.JsonDecodeError((path, "Choice: expected object"))
  }
  let index = match obj.get("index") {
    Some(Number(n, ..)) => n.to_int()
    _ => 0
  }
  let message = match obj.get("message") {
    Some(m) => @json.from_json(m)
    None => raise @json.JsonDecodeError((path, "Choice: missing message"))
  }
  let finish_reason = match obj.get("finish_reason") {
    Some(String(s)) => Some(s)
    _ => None
  }
  { index, message, finish_reason }
}

///|
/// Whether this choice was cut off because it hit the token limit.
pub fn Choice::truncated(self : Choice) -> Bool {
  self.finish_reason is Some("length")
}

///|
/// Whether this choice ended by requesting tool calls.
pub fn Choice::wants_tools(self : Choice) -> Bool {
  self.finish_reason is Some("tool_calls")
}

///|
/// A non-streaming chat completion response.
pub(all) struct ChatResponse {
  id : String
  object : String
  created : Int64
  model : String
  choices : Array[Choice]
  usage : Usage?
  system_fingerprint : String?
} derive(Debug)

///|
pub impl @json.FromJson for ChatResponse with fn from_json(
  json : Json,
  path : @json.JsonPath,
) -> ChatResponse {
  guard json is Object(obj) else {
    raise @json.JsonDecodeError((path, "ChatResponse: expected object"))
  }
  let id = match obj.get("id") {
    Some(String(s)) => s
    _ => ""
  }
  let object = match obj.get("object") {
    Some(String(s)) => s
    _ => "chat.completion"
  }
  let created = match obj.get("created") {
    Some(Number(n, ..)) => n.to_int64()
    _ => 0L
  }
  let model = match obj.get("model") {
    Some(String(s)) => s
    _ => ""
  }
  let choices = match obj.get("choices") {
    Some(Array(_) as c) => @json.from_json(c)
    _ => []
  }
  let usage = match obj.get("usage") {
    Some(Object(_) as u) => Some(@json.from_json(u))
    _ => None
  }
  let system_fingerprint = match obj.get("system_fingerprint") {
    Some(String(s)) => Some(s)
    _ => None
  }
  { id, object, created, model, choices, usage, system_fingerprint }
}

///|
/// The first choice, if the response contains any.
pub fn ChatResponse::first(self : ChatResponse) -> Choice? {
  self.choices.get(0)
}

///|
/// The finish reason of the first choice, if any.
pub fn ChatResponse::finish_reason(self : ChatResponse) -> String? {
  match self.choices.get(0) {
    Some(c) => c.finish_reason
    None => None
  }
}

///|
/// The tool calls requested by the first choice, if any.
pub fn ChatResponse::tool_calls(self : ChatResponse) -> Array[ToolCall] {
  match self.choices.get(0) {
    Some(c) =>
      match c.message.tool_calls {
        Some(tcs) => tcs
        None => []
      }
    None => []
  }
}

///|
/// Whether the first choice requested one or more tool calls.
pub fn ChatResponse::has_tool_calls(self : ChatResponse) -> Bool {
  self.tool_calls().length() > 0
}

///|
/// The total tokens used, or 0 if usage was not reported.
pub fn ChatResponse::total_tokens(self : ChatResponse) -> Int {
  match self.usage {
    Some(u) => u.total_tokens
    None => 0
  }
}

///|
/// Convenience: the text content of the first choice, if any.
pub fn ChatResponse::text(self : ChatResponse) -> String {
  guard self.choices.get(0) is Some(choice) else { return "" }
  match choice.message.content {
    Str(s) => s
    Parts(parts) => {
      let buf = StringBuilder::new()
      for p in parts {
        if p is Text(t) {
          buf.write_string(t)
        }
      }
      buf.to_string()
    }
  }
}