(*---------------------------------------------------------------------------
  Copyright (c) 2026 Invariant Systems. All rights reserved.
  SPDX-License-Identifier: ISC
 ---------------------------------------------------------------------------*)

open Import

let invalid fn message = invalid_arg' "Llm.Error" fn message

let reject_empty fn field value =
  if String.equal "" value then invalid fn (field ^ " must not be empty")

let reject_empty_option fn field = function
  | None -> ()
  | Some value -> reject_empty fn field value

let is_label_char c =
  Char.Ascii.is_lower c || Char.Ascii.is_digit c || Char.equal c '_'

let check_label fn label =
  if String.equal "" label then invalid fn "label must not be empty";
  if not (Char.Ascii.is_lower label.[0]) then
    invalid fn "label must start with a lowercase ASCII letter";
  for i = 1 to String.length label - 1 do
    if not (is_label_char label.[i]) then
      invalid fn
        "label must contain only lowercase ASCII letters, digits, or '_'"
  done

let reserved_kind_labels =
  [
    "cancelled";
    "auth";
    "quota";
    "rate_limited";
    "context_overflow";
    "invalid_request";
    "unsupported";
    "content_policy";
    "transport";
    "timeout";
    "decode";
    "malformed_stream";
    "provider";
  ]

let check_other_kind_label fn label =
  check_label fn label;
  if List.exists (String.equal label) reserved_kind_labels then
    invalid fn ("reserved error kind label: " ^ label)

type phase = Startup | Stream

type kind =
  | Cancelled
  | Auth
  | Quota
  | Rate_limited
  | Context_overflow
  | Invalid_request
  | Unsupported
  | Content_policy
  | Transport
  | Timeout
  | Decode
  | Malformed_stream
  | Provider
  | Other of string

type t = {
  kind : kind;
  phase : phase;
  provider : Provider.t option;
  status : int option;
  request_id : string option;
  redacted_body : string option;
  message : string;
}

let check_kind = function
  | Other label -> check_other_kind_label "v" label
  | Cancelled | Auth | Quota | Rate_limited | Context_overflow | Invalid_request
  | Unsupported | Content_policy | Transport | Timeout | Decode
  | Malformed_stream | Provider ->
      ()

let check_status = function
  | None -> ()
  | Some status ->
      if status < 100 || status > 599 then
        invalid "v" "status must be in the range 100..599"

let v ~kind ?(phase = Startup) ?provider ?status ?request_id ?redacted_body
    message =
  reject_empty "v" "message" message;
  check_kind kind;
  check_status status;
  reject_empty_option "v" "request_id" request_id;
  reject_empty_option "v" "redacted_body" redacted_body;
  { kind; phase; provider; status; request_id; redacted_body; message }

let kind t = t.kind
let phase t = t.phase
let message t = t.message
let provider t = t.provider
let status t = t.status
let request_id t = t.request_id
let redacted_body t = t.redacted_body

let label kind =
  check_kind kind;
  match kind with
  | Cancelled -> "cancelled"
  | Auth -> "auth"
  | Quota -> "quota"
  | Rate_limited -> "rate_limited"
  | Context_overflow -> "context_overflow"
  | Invalid_request -> "invalid_request"
  | Unsupported -> "unsupported"
  | Content_policy -> "content_policy"
  | Transport -> "transport"
  | Timeout -> "timeout"
  | Decode -> "decode"
  | Malformed_stream -> "malformed_stream"
  | Provider -> "provider"
  | Other label -> label

let equal a b = a = b

let pp ppf t =
  match t.provider with
  | None -> Fmt.pf ppf "%s: %s" (label t.kind) t.message
  | Some provider ->
      Fmt.pf ppf "%a:%s: %s" Provider.pp provider (label t.kind) t.message

let phase_label = function Startup -> "startup" | Stream -> "stream"

let phase_of_label = function
  | "startup" -> Some Startup
  | "stream" -> Some Stream
  | _ -> None

let phase_jsont =
  Json.Codec.map ~kind:"LLM error phase"
    ~dec:(fun label ->
      match phase_of_label label with
      | Some phase -> phase
      | None -> decode_error "invalid error phase")
    ~enc:phase_label Json.Codec.string

let kind_of_label = function
  | "cancelled" -> Some Cancelled
  | "auth" -> Some Auth
  | "quota" -> Some Quota
  | "rate_limited" -> Some Rate_limited
  | "context_overflow" -> Some Context_overflow
  | "invalid_request" -> Some Invalid_request
  | "unsupported" -> Some Unsupported
  | "content_policy" -> Some Content_policy
  | "transport" -> Some Transport
  | "timeout" -> Some Timeout
  | "decode" -> Some Decode
  | "malformed_stream" -> Some Malformed_stream
  | "provider" -> Some Provider
  | label -> (
      try
        check_other_kind_label "kind_of_label" label;
        Some (Other label)
      with Invalid_argument _ -> None)

let kind_jsont =
  Json.Codec.map ~kind:"LLM error kind"
    ~dec:(fun label ->
      match kind_of_label label with
      | Some kind -> kind
      | None -> decode_error "invalid error kind")
    ~enc:label Json.Codec.string

let jsont =
  let v kind phase provider status request_id redacted_body message =
    decode_invalid_arg (fun () ->
        v ~kind ~phase ?provider ?status ?request_id ?redacted_body message)
  in
  Json.Codec.Object.map ~kind:"LLM error" v
  |> Json.Codec.Object.member "kind" kind_jsont ~enc:kind
  |> Json.Codec.Object.member "phase" phase_jsont ~enc:phase
  |> Json.Codec.Object.opt_member "provider" Provider.jsont ~enc:provider
  |> Json.Codec.Object.opt_member "status" Json.Codec.int ~enc:status
  |> Json.Codec.Object.opt_member "request_id" Json.Codec.string ~enc:request_id
  |> Json.Codec.Object.opt_member "redacted_body" Json.Codec.string
       ~enc:redacted_body
  |> Json.Codec.Object.member "message" Json.Codec.string ~enc:message
  |> Json.Codec.Object.error_unknown |> Json.Codec.Object.seal