module Encoding = struct
  type t = Utf8 | Utf16le | Utf16be | Utf32le | Utf32be | Latin1

  let name = function
    | Utf8 -> "UTF-8"
    | Utf16le -> "UTF-16LE"
    | Utf16be -> "UTF-16BE"
    | Utf32le -> "UTF-32LE"
    | Utf32be -> "UTF-32BE"
    | Latin1 -> "Latin-1"

  external detect_mask : string -> int = "ocaml_simdutf_detect"

  let of_mask mask =
    List.filter_map
      (fun (bit, encoding) ->
        if mask land bit <> 0 then Some encoding else None)
      [ (1, Utf8); (2, Utf16le); (4, Utf16be); (8, Utf32le); (16, Utf32be) ]

  let detect input = of_mask (detect_mask input)
end

type error_code =
  | Header_bits
  | Too_short
  | Too_long
  | Overlong
  | Too_large
  | Surrogate
  | Invalid_base64_character
  | Base64_input_remainder
  | Base64_extra_bits
  | Output_buffer_too_small
  | Other
  | Incomplete_code_unit

type error = { code : error_code; offset : int; message : string }

exception Error of error

let make_error code offset =
  let code, message =
    match code with
    | 1 -> (Header_bits, "Invalid UTF-8 leading byte")
    | 2 -> (Too_short, "Truncated UTF-8 sequence")
    | 3 -> (Too_long, "Unexpected UTF-8 continuation byte")
    | 4 -> (Overlong, "Overlong UTF-8 sequence")
    | 5 -> (Too_large, "Code point is out of range for the encoding")
    | 6 -> (Surrogate, "Invalid or unrepresentable surrogate")
    | 7 -> (Invalid_base64_character, "Invalid Base64 character or padding")
    | 8 -> (Base64_input_remainder, "Incomplete Base64 group or padding")
    | 9 -> (Base64_extra_bits, "Non-zero unused Base64 bits")
    | 10 -> (Output_buffer_too_small, "Output buffer is too small")
    | 12 -> (Incomplete_code_unit, "Incomplete UTF-16 or UTF-32 code unit")
    | _ -> (Other, "simdutf error")
  in
  { code; offset; message }

let pp_error ppf e = Format.fprintf ppf "%s at byte %d" e.message e.offset

(* Stubs return one immediate integer: a count >= 0 on success, or
   -(code + 16 * offset). Codes 13 to 15 report failures that the stubs cannot
   raise themselves. Keep in sync with simdutf_stubs.cpp. *)
let status_error r =
  let r = -r in
  match r land 15 with
  | 13 -> invalid_arg "Simdutf: overlapping input and output slices"
  | 14 -> raise Out_of_memory
  | 15 -> invalid_arg "Simdutf: output exceeds maximum buffer length"
  | code -> Result.Error (make_error code (r lsr 4))

let checked_unit r = if r >= 0 then Ok () else status_error r
let checked_count r = if r >= 0 then Ok r else status_error r

let checked_bool r =
  if r >= 0 then r <> 0
  else (
    ignore (status_error r : (unit, error) result);
    false)

(* An output of [capacity] bytes, of which the stub wrote [r]. *)
let checked_output r capacity output =
  if r = capacity then Ok (Bytes.unsafe_to_string output)
  else if r >= 0 then Ok (Bytes.sub_string output 0 r)
  else status_error r

let get_exn = function Ok v -> v | Result.Error e -> raise (Error e)

(* The default of every optional length. A defaulted optional argument avoids
   allocating [Some len] at call sites, once the wrapper that the compiler
   generates for it is inlined. An explicit [~len:min_int] thus means the
   remaining bytes; any other negative length is rejected. *)
let rest = min_int

let check_slice name size off len =
  if off < 0 || off > size then invalid_arg (name ^ ": invalid offset");
  let len = if len = rest then size - off else len in
  if len < 0 || len > size - off then invalid_arg (name ^ ": invalid length");
  len

let string_slice s off len = check_slice "Simdutf" (String.length s) off len
let bytes_slice b off len = check_slice "Simdutf" (Bytes.length b) off len

external upstream_version : unit -> string = "ocaml_simdutf_version"
external implementation : unit -> string = "ocaml_simdutf_implementation"

let version = upstream_version ()

external validate_raw : Encoding.t -> string -> int -> int -> int
  = "ocaml_simdutf_validate"
[@@noalloc]

external is_valid_raw : Encoding.t -> string -> int -> int -> int
  = "ocaml_simdutf_is_valid"
[@@noalloc]

external ascii_raw : string -> int -> int -> int
  = "ocaml_simdutf_validate_ascii"
[@@noalloc]

external is_ascii_raw : string -> int -> int -> bool = "ocaml_simdutf_is_ascii"
[@@noalloc]

external count_raw : Encoding.t -> string -> int -> int -> int
  = "ocaml_simdutf_count"
[@@noalloc]

let validate ?(encoding = Encoding.Utf8) ?(off = 0) ?(len = rest) input =
  let len = string_slice input off len in
  checked_unit (validate_raw encoding input off len)

let is_valid ?(encoding = Encoding.Utf8) ?(off = 0) ?(len = rest) input =
  let len = string_slice input off len in
  checked_bool (is_valid_raw encoding input off len)

let validate_ascii ?(off = 0) ?(len = rest) input =
  let len = string_slice input off len in
  checked_unit (ascii_raw input off len)

let is_ascii ?(off = 0) ?(len = rest) input =
  let len = string_slice input off len in
  is_ascii_raw input off len

let count_codepoints ?(encoding = Encoding.Utf8) ?(off = 0) ?(len = rest) input
    =
  let len = string_slice input off len in
  checked_count (count_raw encoding input off len)

(* Search stubs return the match offset within the slice, or the length of
   the searched part when there is none. *)
let found off searched r = if r = searched then -1 else off + r

let check_code_unit name (encoding : Encoding.t) code_unit =
  (match encoding with
  | Utf16le | Utf16be -> ()
  | _ -> invalid_arg (name ^ ": encoding must be UTF-16LE or UTF-16BE"));
  if code_unit < 0 || code_unit > 0xffff then
    invalid_arg (name ^ ": code unit out of range")

external find_raw : string -> int -> int -> char -> int = "ocaml_simdutf_find"
[@@noalloc]

external find_code_unit_raw : Encoding.t -> int -> string -> int -> int -> int
  = "ocaml_simdutf_find_code_unit"
[@@noalloc]

let find ?(off = 0) ?(len = rest) input byte =
  let len = string_slice input off len in
  found off len (find_raw input off len byte)

let find_utf16 ~encoding ?(off = 0) ?(len = rest) input code_unit =
  let len = string_slice input off len in
  check_code_unit "Simdutf.find_utf16" encoding code_unit;
  found off
    (len land lnot 1)
    (find_code_unit_raw encoding code_unit input off len)

external transcode_into_raw :
  Encoding.t -> Encoding.t -> string -> int -> int -> bytes -> int -> int -> int
  = "ocaml_simdutf_transcode_into_bytes_bytecode"
    "ocaml_simdutf_transcode_into_bytes"
[@@noalloc]

let unit_size = function
  | Encoding.Utf8 | Latin1 -> 1
  | Utf16le | Utf16be -> 2
  | Utf32le | Utf32be -> 4

(* An upper bound on the output size, so that conversion needs one pass. *)
let transcode_capacity src dst len =
  if src = dst then len
  else
    let units = len / unit_size src in
    let factor =
      match (dst : Encoding.t) with
      | Utf8 -> (
          match src with Latin1 -> 2 | Utf16le | Utf16be -> 3 | _ -> 4)
      | Utf16le | Utf16be -> if unit_size src = 4 then 4 else 2
      | Utf32le | Utf32be | Latin1 -> unit_size dst
    in
    if units > Sys.max_string_length / factor then
      invalid_arg "Simdutf: output exceeds maximum string length";
    units * factor

let transcode ~src ~dst ?(off = 0) ?(len = rest) input =
  let len = string_slice input off len in
  let capacity = transcode_capacity src dst len in
  let output = Bytes.create capacity in
  checked_output
    (transcode_into_raw src dst input off len output 0 capacity)
    capacity output

let transcode_exn ~src ~dst ?off ?len input =
  get_exn (transcode ~src ~dst ?off ?len input)

module Base64 = struct
  type alphabet = Standard | Url
  type last_chunk = Loose | Strict

  external encode_raw : int -> string -> string = "ocaml_simdutf_base64_encode"

  external decoded_length_raw : string -> int -> int -> int
    = "ocaml_simdutf_base64_decoded_length"
  [@@noalloc]

  external decode_into_raw :
    int -> string -> int -> int -> bytes -> int -> int -> int
    = "ocaml_simdutf_base64_decode_into_bytecode"
      "ocaml_simdutf_base64_decode_into"
  [@@noalloc]

  let encoding_options alphabet padding =
    let default_padding = alphabet = Standard in
    let padding = Option.value padding ~default:default_padding in
    let options = if alphabet = Standard then 0 else 1 in
    if padding = default_padding then options else options lor 2

  (* Bit 2, set by callers that sized the output with [decoded_length] on the
     same immutable input, skips the capacity check. *)
  let decode_mode alphabet last_chunk =
    (match alphabet with Standard -> 0 | Url -> 1)
    lor match last_chunk with Loose -> 0 | Strict -> 2

  let trusted_capacity = 4

  let encode ?(alphabet = Standard) ?padding input =
    encode_raw (encoding_options alphabet padding) input

  let decoded_length ?(off = 0) ?(len = rest) input =
    let len = string_slice input off len in
    decoded_length_raw input off len

  let decode_into ?(alphabet = Standard) ?(last_chunk = Loose) ?(off = 0)
      ?(len = rest) ?(output_off = 0) ?(output_len = rest) ~output input =
    let len = string_slice input off len in
    let output_len = bytes_slice output output_off output_len in
    checked_count
      (decode_into_raw
         (decode_mode alphabet last_chunk)
         input off len output output_off output_len)

  let decode ?(alphabet = Standard) ?(last_chunk = Loose) ?(off = 0)
      ?(len = rest) input =
    let len = string_slice input off len in
    let capacity = decoded_length_raw input off len in
    let output = Bytes.create capacity in
    checked_output
      (decode_into_raw
         (decode_mode alphabet last_chunk lor trusted_capacity)
         input off len output 0 capacity)
      capacity output

  let decode_exn ?alphabet ?last_chunk ?off ?len input =
    get_exn (decode ?alphabet ?last_chunk ?off ?len input)
end

module Bigstring = struct
  type t =
    (char, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t

  let create len = Bigarray.Array1.create Bigarray.char Bigarray.c_layout len

  let slice b off len =
    check_slice "Simdutf.Bigstring" (Bigarray.Array1.dim b) off len

  external blit_from_string : string -> int -> t -> int -> int -> unit
    = "ocaml_simdutf_blit_string_to_bigstring"
  [@@noalloc]

  external blit_to_bytes : t -> int -> bytes -> int -> int -> unit
    = "ocaml_simdutf_blit_bigstring_to_bytes"
  [@@noalloc]

  let of_string s =
    let len = String.length s in
    let b = create len in
    blit_from_string s 0 b 0 len;
    b

  let sub_string ?(off = 0) ?(len = rest) b =
    let len = slice b off len in
    let s = Bytes.create len in
    blit_to_bytes b off s 0 len;
    Bytes.unsafe_to_string s

  let to_string b = sub_string b

  (* Inputs of at least this many bytes use the [_blocking] stubs, which
     release the runtime lock. Below it, the [@@noalloc] stubs avoid the
     cost of releasing and reacquiring the lock. *)
  let blocking = 65536

  external validate_raw : Encoding.t -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_validate"
  [@@noalloc]

  external validate_blocking : Encoding.t -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_validate_blocking"

  external is_valid_raw : Encoding.t -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_is_valid"
  [@@noalloc]

  external is_valid_blocking : Encoding.t -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_is_valid_blocking"

  external count_raw : Encoding.t -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_count"
  [@@noalloc]

  external count_blocking : Encoding.t -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_count_blocking"

  let validate ?(encoding = Encoding.Utf8) ?(off = 0) ?(len = rest) b =
    let len = slice b off len in
    checked_unit
      (if len < blocking then validate_raw encoding b off len
       else validate_blocking encoding b off len)

  let is_valid ?(encoding = Encoding.Utf8) ?(off = 0) ?(len = rest) b =
    let len = slice b off len in
    checked_bool
      (if len < blocking then is_valid_raw encoding b off len
       else is_valid_blocking encoding b off len)

  let count_codepoints ?(encoding = Encoding.Utf8) ?(off = 0) ?(len = rest) b =
    let len = slice b off len in
    checked_count
      (if len < blocking then count_raw encoding b off len
       else count_blocking encoding b off len)

  external ascii_raw : t -> int -> int -> int
    = "ocaml_simdutf_bigstring_validate_ascii"
  [@@noalloc]

  external ascii_blocking : t -> int -> int -> int
    = "ocaml_simdutf_bigstring_validate_ascii_blocking"

  external is_ascii_raw : t -> int -> int -> int
    = "ocaml_simdutf_bigstring_is_ascii"
  [@@noalloc]

  external is_ascii_blocking : t -> int -> int -> int
    = "ocaml_simdutf_bigstring_is_ascii_blocking"

  let validate_ascii ?(off = 0) ?(len = rest) b =
    let len = slice b off len in
    checked_unit
      (if len < blocking then ascii_raw b off len else ascii_blocking b off len)

  let is_ascii ?(off = 0) ?(len = rest) b =
    let len = slice b off len in
    checked_bool
      (if len < blocking then is_ascii_raw b off len
       else is_ascii_blocking b off len)

  external transcoded_length_raw :
    Encoding.t -> Encoding.t -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_transcoded_length"
  [@@noalloc]

  external transcoded_length_blocking :
    Encoding.t -> Encoding.t -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_transcoded_length_blocking"

  external transcode_into_raw :
    Encoding.t -> Encoding.t -> t -> int -> int -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_transcode_into_bytecode"
      "ocaml_simdutf_bigstring_transcode_into"
  [@@noalloc]

  external transcode_into_blocking :
    Encoding.t -> Encoding.t -> t -> int -> int -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_transcode_into_blocking_bytecode"
      "ocaml_simdutf_bigstring_transcode_into_blocking"

  let transcoded_length ~src ~dst ?(off = 0) ?(len = rest) b =
    let len = slice b off len in
    checked_count
      (if len < blocking then transcoded_length_raw src dst b off len
       else transcoded_length_blocking src dst b off len)

  let transcode_into ~src ~dst ?(off = 0) ?(len = rest) ?(output_off = 0)
      ?(output_len = rest) ~output b =
    let len = slice b off len in
    let output_len = slice output output_off output_len in
    checked_count
      (if len < blocking then
         transcode_into_raw src dst b off len output output_off output_len
       else
         transcode_into_blocking src dst b off len output output_off output_len)

  let transcode ~src ~dst ?(off = 0) ?(len = rest) b =
    let len = slice b off len in
    Result.bind (transcoded_length ~src ~dst ~off ~len b) (fun capacity ->
        let output = create capacity in
        Result.map
          (fun _ -> output)
          (transcode_into ~src ~dst ~off ~len ~output b))

  let transcode_exn ~src ~dst ?off ?len b =
    get_exn (transcode ~src ~dst ?off ?len b)

  external find_raw : t -> int -> int -> char -> int
    = "ocaml_simdutf_bigstring_find"
  [@@noalloc]

  external find_code_unit_raw : Encoding.t -> int -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_find_code_unit"
  [@@noalloc]

  external find_code_unit_blocking : Encoding.t -> int -> t -> int -> int -> int
    = "ocaml_simdutf_bigstring_find_code_unit_blocking"

  (* Never releases the runtime lock, unlike the operations below. *)
  let find ?(off = 0) ?(len = rest) b byte =
    let len = slice b off len in
    found off len (find_raw b off len byte)

  let find_utf16 ~encoding ?(off = 0) ?(len = rest) b code_unit =
    let len = slice b off len in
    check_code_unit "Simdutf.Bigstring.find_utf16" encoding code_unit;
    found off
      (len land lnot 1)
      (if len < blocking then find_code_unit_raw encoding code_unit b off len
       else find_code_unit_blocking encoding code_unit b off len)

  module Encoding = struct
    external detect_raw : t -> int -> int -> int
      = "ocaml_simdutf_bigstring_detect"

    let detect ?(off = 0) ?(len = rest) b =
      let len = slice b off len in
      Encoding.of_mask (detect_raw b off len)
  end

  module Base64 = struct
    type alphabet = Base64.alphabet = Standard | Url
    type last_chunk = Base64.last_chunk = Loose | Strict

    external encoded_length_raw : int -> int -> int
      = "ocaml_simdutf_bigstring_base64_encoded_length"

    external decoded_length_raw : t -> int -> int -> int
      = "ocaml_simdutf_bigstring_base64_decoded_length"
    [@@noalloc]

    external decoded_length_blocking : t -> int -> int -> int
      = "ocaml_simdutf_bigstring_base64_decoded_length_blocking"

    external encode_into_raw : int -> t -> int -> int -> t -> int -> int -> int
      = "ocaml_simdutf_bigstring_base64_encode_into_bytecode"
        "ocaml_simdutf_bigstring_base64_encode_into"
    [@@noalloc]

    external encode_into_blocking :
      int -> t -> int -> int -> t -> int -> int -> int
      = "ocaml_simdutf_bigstring_base64_encode_into_blocking_bytecode"
        "ocaml_simdutf_bigstring_base64_encode_into_blocking"

    external decode_into_raw : int -> t -> int -> int -> t -> int -> int -> int
      = "ocaml_simdutf_bigstring_base64_decode_into_bytecode"
        "ocaml_simdutf_bigstring_base64_decode_into"
    [@@noalloc]

    external decode_into_blocking :
      int -> t -> int -> int -> t -> int -> int -> int
      = "ocaml_simdutf_bigstring_base64_decode_into_blocking_bytecode"
        "ocaml_simdutf_bigstring_base64_decode_into_blocking"

    external decode_to_bytes_raw :
      int -> t -> int -> int -> bytes -> int -> int -> int
      = "ocaml_simdutf_bigstring_base64_decode_to_bytes_bytecode"
        "ocaml_simdutf_bigstring_base64_decode_to_bytes"
    [@@noalloc]

    let encoded_length ?(alphabet = Standard) ?padding len =
      encoded_length_raw (Base64.encoding_options alphabet padding) len

    let decoded_length ?(off = 0) ?(len = rest) b =
      let len = slice b off len in
      if len < blocking then decoded_length_raw b off len
      else decoded_length_blocking b off len

    let encode_into ?(alphabet = Standard) ?padding ?(off = 0) ?(len = rest)
        ?(output_off = 0) ?(output_len = rest) ~output b =
      let len = slice b off len in
      let output_len = slice output output_off output_len in
      let options = Base64.encoding_options alphabet padding in
      checked_count
        (if len < blocking then
           encode_into_raw options b off len output output_off output_len
         else
           encode_into_blocking options b off len output output_off output_len)

    let decode_into ?(alphabet = Standard) ?(last_chunk = Loose) ?(off = 0)
        ?(len = rest) ?(output_off = 0) ?(output_len = rest) ~output b =
      let len = slice b off len in
      let output_len = slice output output_off output_len in
      let mode = Base64.decode_mode alphabet last_chunk in
      checked_count
        (if len < blocking then
           decode_into_raw mode b off len output output_off output_len
         else decode_into_blocking mode b off len output output_off output_len)

    let decode_to_bytes ?(alphabet = Standard) ?(last_chunk = Loose) ?(off = 0)
        ?(len = rest) ?(output_off = 0) ?(output_len = rest) ~output b =
      let len = slice b off len in
      let output_len = bytes_slice output output_off output_len in
      checked_count
        (decode_to_bytes_raw
           (Base64.decode_mode alphabet last_chunk)
           b off len output output_off output_len)

    let encode ?(alphabet = Standard) ?padding ?(off = 0) ?(len = rest) b =
      let len = slice b off len in
      let output = create (encoded_length ~alphabet ?padding len) in
      ignore (get_exn (encode_into ~alphabet ?padding ~off ~len ~output b));
      output

    let decode ?(alphabet = Standard) ?(last_chunk = Loose) ?(off = 0)
        ?(len = rest) b =
      let len = slice b off len in
      let output = create (decoded_length ~off ~len b) in
      Result.map
        (fun written ->
          if written = Bigarray.Array1.dim output then output
          else Bigarray.Array1.sub output 0 written)
        (decode_into ~alphabet ~last_chunk ~off ~len ~output b)

    let decode_exn ?alphabet ?last_chunk ?off ?len b =
      get_exn (decode ?alphabet ?last_chunk ?off ?len b)
  end
end