open Prelude
open Option.Let_syntax

type t = string

(* Modulo-97 arithmetic. Prevents us from having to use Zarith here. *)
module M97 : sig
  type t = private int

  val of_int : int -> t
  val lt : t -> t -> bool
  val equal : t -> t -> bool
  val ( * ) : t -> t -> t
  val ( + ) : t -> t -> t
  val ( ~$ ) : int -> t
end = struct
  type t = int

  let of_int x = x % 97
  let equal = Int.( = )
  let lt = Int.( < )
  let ( * ) x y = x * y % 97
  let ( + ) x y = (x + y) % 97
  let ( ~$ ) = of_int
end

let m97_of_alnum c =
  let v = Char.to_int c in
  if Char.is_digit c then Some (M97.of_int (v - Char.to_int '0'))
  else if Char.is_alpha c then
    if Char.is_lowercase c then Some (M97.of_int (v - Char.to_int 'a' + 10))
    else Some (M97.of_int (v - Char.to_int 'A' + 10))
  else None

let m97_of_digit c =
  match m97_of_alnum c with Some v when M97.(lt v ~$10) -> Some v | _ -> None

let m97_of_alpha c =
  match m97_of_alnum c with
  | Some v when not M97.(lt v ~$10) -> Some v
  | _ -> None

let string_fold_option ~(init : 'a) ~(f : 'a -> char -> 'a option) s =
  let rec go i (acc : 'a) : 'a option =
    if i >= String.length s then Some acc
    else Option.(f acc (String.unsafe_get s i) >>= go (i + 1))
  in
  go 0 init

let m97_of_iban s =
  string_fold_option s ~init:`In_country1 ~f:(fun st c ->
      match st with
      | `In_country1 ->
          let%map co1 = m97_of_alpha c in
          `In_country2 co1
      | `In_country2 co1 ->
          let%map co2 = m97_of_alpha c in
          `In_check1 M97.((co1 * ~$100) + co2)
      | `In_check1 co ->
          let%map ch1 = m97_of_digit c in
          `In_check2 (co, ch1)
      | `In_check2 (co, ch1) ->
          let%map ch2 = m97_of_digit c in
          `In_bban M97.(co, (ch1 * ~$10) + ch2, ~$0)
      | `In_bban (co, ch, bban) ->
          let%map v = m97_of_alnum c in
          let bban' =
            M97.(if lt v ~$10 then (bban * ~$10) + v else (bban * ~$100) + v)
          in
          `In_bban (co, ch, bban'))
  |> function
  | Some (`In_bban (co, ch, bban)) ->
      Some M97.((bban * ~$1000000) + (co * ~$100) + ch)
  | _ -> None

let check_iban s =
  String.length s <= 34 && Option.exists (m97_of_iban s) ~f:M97.(equal ~$1)

let make s : t option = if check_iban s then Some s else None
let to_string = Fn.id

let of_string s =
  match make s with
  | Some iban -> iban
  | None -> Printf.failwithf "Iban.of_string: %S" s ()

let sexp_of_t iban = Sexp.Atom iban
let equal = String.equal