(*---------------------------------------------------------------------------
   Copyright (c) 2025 Anil Madhavapeddy. All rights reserved.
   SPDX-License-Identifier: MIT
  ---------------------------------------------------------------------------*)

type invalid_length = { length: int; message: string }
type invalid_character = { char: char; message: string }
type invalid_checksum = { checksum: string; message: string }
type checksum_mismatch = { expected: int64; got: int64; identifier: string }

type decode_error =
  | Invalid_length of invalid_length
  | Invalid_character of invalid_character
  | Invalid_checksum of invalid_checksum
  | Checksum_mismatch of checksum_mismatch

exception Decode_error of decode_error

let pp_invalid_length fmt { length; message } =
  Format.fprintf fmt "Invalid_length: length=%d, %s" length message

let pp_invalid_character fmt { char; message } =
  Format.fprintf fmt "Invalid_character: char='%c', %s" char message

let pp_invalid_checksum fmt { checksum; message } =
  Format.fprintf fmt "Invalid_checksum: checksum=%s, %s" checksum message

let pp_checksum_mismatch fmt { expected; got; identifier } =
  Format.fprintf fmt "Checksum_mismatch: expected=%Ld, got=%Ld, identifier=%s"
    expected got identifier

let pp_decode_error fmt = function
  | Invalid_length e -> pp_invalid_length fmt e
  | Invalid_character e -> pp_invalid_character fmt e
  | Invalid_checksum e -> pp_invalid_checksum fmt e
  | Checksum_mismatch e -> pp_checksum_mismatch fmt e

let encoding_chars = "0123456789abcdefghjkmnpqrstvwxyz"

let generate_checksum number =
  Int64.(sub (add (sub 97L (rem (mul 100L number) 97L)) 1L) 0L)

let validate number ~checksum =
  Int64.equal checksum (generate_checksum number)

let normalize str =
  let len = String.length str in
  let buf = Bytes.create len in
  let rec process i j =
    if i >= len then Bytes.sub_string buf 0 j
    else
      let c = String.get str i in
      let c_lower = Char.lowercase_ascii c in
      match c_lower with
      | '-' -> process (i + 1) j
      | 'i' | 'l' -> Bytes.set buf j '1'; process (i + 1) (j + 1)
      | 'o' -> Bytes.set buf j '0'; process (i + 1) (j + 1)
      | _ -> Bytes.set buf j c_lower; process (i + 1) (j + 1)
  in
  process 0 0

let encode ?(split_every=0) ?(min_length=0) ?(checksum=false) number =
  let original_number = number in

  (* Build base32 encoding *)
  let rec build_encoding acc n =
    if Int64.equal n 0L then acc
    else
      let remainder = Int64.to_int (Int64.rem n 32L) in
      let n' = Int64.div n 32L in
      build_encoding (encoding_chars.[remainder] :: acc) n'
  in

  let encoded_list =
    if Int64.equal number 0L then ['0']
    else build_encoding [] number
  in

  let encoded_str = String.concat "" (List.map (String.make 1) encoded_list) in

  (* Adjust min_length if checksum is enabled *)
  let adjusted_length =
    if checksum && min_length > 2 then min_length - 2
    else min_length
  in

  (* Pad with zeros if needed *)
  let padded =
    if adjusted_length > 0 && String.length encoded_str < adjusted_length then
      String.make (adjusted_length - String.length encoded_str) '0' ^ encoded_str
    else
      encoded_str
  in

  (* Add checksum *)
  let with_checksum =
    if checksum then
      let cs = generate_checksum original_number in
      padded ^ Printf.sprintf "%02Ld" cs
    else
      padded
  in

  (* Split if requested *)
  if split_every > 0 then
    let len = String.length with_checksum in
    let num_splits = (len + split_every - 1) / split_every in
    let splits = Array.make num_splits "" in
    for i = 0 to num_splits - 1 do
      let start = i * split_every in
      let chunk_len = min split_every (len - start) in
      splits.(i) <- String.sub with_checksum start chunk_len
    done;
    String.concat "-" (Array.to_list splits)
  else
    with_checksum

let decode ?(checksum=false) str =
  let encoded = normalize str in

  let (encoded_part, checksum_value) =
    if checksum then begin
      if String.length encoded < 3 then
        raise (Decode_error (Invalid_checksum {
          checksum = encoded;
          message = "encoded string too short for checksum"
        }));

      let cs_str = String.sub encoded (String.length encoded - 2) 2 in
      let cs =
        try Int64.of_string cs_str
        with Failure _ ->
          raise (Decode_error (Invalid_checksum {
            checksum = cs_str;
            message = "invalid checksum format"
          }))
      in
      (String.sub encoded 0 (String.length encoded - 2), Some cs)
    end else
      (encoded, None)
  in

  (* Decode base32 *)
  let number = ref 0L in
  String.iter (fun c ->
    number := Int64.mul !number 32L;
    match String.index_opt encoding_chars c with
    | Some pos -> number := Int64.add !number (Int64.of_int pos)
    | None ->
        raise (Decode_error (Invalid_character {
          char = c;
          message = Printf.sprintf "character '%c' not in base32 alphabet" c
        }))
  ) encoded_part;

  (* Validate checksum if present *)
  (match checksum_value with
  | Some cs ->
      if not (validate !number ~checksum:cs) then
        raise (Decode_error (Checksum_mismatch {
          expected = generate_checksum !number;
          got = cs;
          identifier = str
        }))
  | None -> ());

  !number

let generate ~length ?(split_every=0) ?(checksum=false) ?(rng=Random.float) () =
  if checksum && length < 3 then
    raise (Decode_error (Invalid_length {
      length;
      message = "length must be >= 3 if checksum is enabled"
    }));

  let adjusted_length = if checksum then length - 2 else length in

  (* Generate random number between 0 and 32^length *)
  let max_val = 32.0 ** float_of_int adjusted_length in
  let random_num = Int64.of_float (rng max_val) in

  encode ~split_every ~min_length:adjusted_length ~checksum random_num
