let version = "0.2.0"
let default_url = "https://cep.is"
let max_batch = 500

type geo = { lat : float; lon : float; raio_m : int; precisao : string }

let is_precise g = String.equal g.precisao "logradouro"

type address = {
  cep : string;
  logradouro : string;
  complemento : string;
  unidade : string;
  bairro : string;
  localidade : string;
  uf : string;
  estado : string;
  regiao : string;
  ibge : string;
  gia : string;
  ddd : string;
  siafi : string;
  geo : geo option;
  raw : (string * Yojson.Safe.t) list;
}

let formatted_cep a =
  if String.length a.cep = 8 then
    String.sub a.cep 0 5 ^ "-" ^ String.sub a.cep 5 3
  else a.cep

type health = {
  status : string;
  dne : string;
  ceps : int;
  geo : int;
  base_gerada_em : string;
}

let is_ok h = String.equal h.status "ok"

type faixa = {
  localidade : string;
  uf : string;
  estado : string;
  regiao : string;
  ibge : string;
  cep_ini : string;
  cep_fim : string;
  ceps : int;
}

type error =
  | Invalid_cep of string
  | Cep_not_found of { message : string; faixa : faixa option }
  | Rate_limited of { message : string; retry_after : float option }
  | Server_error of { message : string; status : int }
  | Network_error of string
  | Invalid_response of string

let error_message = function
  | Invalid_cep m | Network_error m | Invalid_response m -> m
  | Cep_not_found { message; _ }
  | Rate_limited { message; _ }
  | Server_error { message; _ } ->
      message

let string_of_error = function
  | Invalid_cep m -> "invalid CEP: " ^ m
  | Cep_not_found { message; faixa = None } -> "CEP not found: " ^ message
  | Cep_not_found { message; faixa = Some f } ->
      Printf.sprintf "CEP not found: %s (that number falls in %s/%s)" message
        f.localidade f.uf
  | Rate_limited { message; retry_after = Some s } ->
      Printf.sprintf "rate limited (retry after %gs): %s" s message
  | Rate_limited { message; retry_after = None } -> "rate limited: " ^ message
  | Server_error { message; status } ->
      Printf.sprintf "server error (%d): %s" status message
  | Network_error m -> "network error: " ^ m
  | Invalid_response m -> "invalid response: " ^ m

let is_digit c = c >= '0' && c <= '9'

let normalize_cep value =
  let buf = Buffer.create 8 in
  String.iter
    (fun c -> match c with '-' | '.' | ' ' -> () | c -> Buffer.add_char buf c)
    (String.trim value);
  let cleaned = Buffer.contents buf in
  let ok =
    String.length cleaned = 8 && String.for_all (fun c -> is_digit c) cleaned
  in
  if ok then Ok cleaned
  else Error (Invalid_cep (Printf.sprintf "CEP must be 8 digits, got %S" value))

let json_to_string = function
  | `String s -> s
  | `Int i -> string_of_int i
  | `Intlit s -> s
  | `Float f ->
      if Float.is_integer f then Printf.sprintf "%.0f" f
      else Printf.sprintf "%g" f
  | `Bool b -> string_of_bool b
  | _ -> ""

let json_to_float = function
  | `Int i -> Some (float_of_int i)
  | `Float f -> Some f
  | `Intlit s | `String s -> float_of_string_opt s
  | _ -> None

let field fields key =
  match List.assoc_opt key fields with Some v -> v | None -> `Null

let string_field fields key = json_to_string (field fields key)

let int_field fields key =
  match field fields key with
  | `Int i -> i
  | `Float f -> int_of_float f
  | `Intlit s | `String s -> (
      match int_of_string_opt s with Some i -> i | None -> 0)
  | _ -> 0

let geo_of_json = function
  | `Assoc fields -> (
      match
        ( json_to_float (field fields "lat"),
          json_to_float (field fields "lon"),
          json_to_float (field fields "raio_m"),
          field fields "precisao" )
      with
      | Some lat, Some lon, Some raio_m, ((`String _ | `Int _ | `Float _) as p)
        ->
          Some
            {
              lat;
              lon;
              raio_m = int_of_float raio_m;
              precisao = json_to_string p;
            }
      | _ -> None)
  | _ -> None

let faixa_of_json : Yojson.Safe.t -> faixa option = function
  | `Assoc fields -> (
      match string_field fields "localidade" with
      | "" -> None
      | localidade ->
          Some
            {
              localidade;
              uf = string_field fields "uf";
              estado = string_field fields "estado";
              regiao = string_field fields "regiao";
              ibge = string_field fields "ibge";
              cep_ini = string_field fields "cep_ini";
              cep_fim = string_field fields "cep_fim";
              ceps = int_field fields "ceps";
            })
  | _ -> None

let address_of_fields fields =
  {
    cep = string_field fields "cep";
    logradouro = string_field fields "logradouro";
    complemento = string_field fields "complemento";
    unidade = string_field fields "unidade";
    bairro = string_field fields "bairro";
    localidade = string_field fields "localidade";
    uf = string_field fields "uf";
    estado = string_field fields "estado";
    regiao = string_field fields "regiao";
    ibge = string_field fields "ibge";
    gia = string_field fields "gia";
    ddd = string_field fields "ddd";
    siafi = string_field fields "siafi";
    geo = geo_of_json (field fields "geo");
    raw = fields;
  }

let health_of_fields fields =
  {
    status = string_field fields "status";
    dne = string_field fields "dne";
    ceps = int_field fields "ceps";
    geo = int_field fields "geo";
    base_gerada_em = string_field fields "base_gerada_em";
  }

let json_kind = function
  | `Assoc _ -> "object"
  | `List _ -> "array"
  | `String _ -> "string"
  | `Int _ | `Intlit _ -> "int"
  | `Float _ -> "float"
  | `Bool _ -> "bool"
  | `Null -> "null"
  | _ -> "value"

let truncate_body body =
  if String.length body <= 120 then body else String.sub body 0 120 ^ "..."

let json_object_of_body body =
  match Yojson.Safe.from_string body with
  | `Assoc fields -> Ok fields
  | other ->
      Error
        (Invalid_response
           (Printf.sprintf "expected a JSON object, got %s" (json_kind other)))
  | exception _ ->
      Error
        (Invalid_response
           (Printf.sprintf
              "response is not valid JSON — something between you and the API \
               is answering instead of it: %S"
              (truncate_body body)))

let api_message body fallback =
  match Yojson.Safe.from_string body with
  | `Assoc fields -> (
      match field fields "mensagem" with
      | `String m when m <> "" -> m
      | _ -> fallback)
  | _ -> fallback
  | exception _ -> fallback

let faixa_of_body body =
  match Yojson.Safe.from_string body with
  | `Assoc fields -> faixa_of_json (field fields "faixa")
  | _ -> None
  | exception _ -> None

let months =
  [|
    "jan";
    "feb";
    "mar";
    "apr";
    "may";
    "jun";
    "jul";
    "aug";
    "sep";
    "oct";
    "nov";
    "dec";
  |]

let month_index name =
  let name = String.lowercase_ascii name in
  let name = if String.length name > 3 then String.sub name 0 3 else name in
  let found = ref None in
  Array.iteri
    (fun i m -> if String.equal m name then found := Some (i + 1))
    months;
  !found

let days_from_civil y m d =
  let y = if m <= 2 then y - 1 else y in
  let era = (if y >= 0 then y else y - 399) / 400 in
  let yoe = y - (era * 400) in
  let mp = (m + 9) mod 12 in
  let doy = (((153 * mp) + 2) / 5) + d - 1 in
  let doe = (yoe * 365) + (yoe / 4) - (yoe / 100) + doy in
  (era * 146097) + doe - 719468

let words s =
  String.split_on_char ' ' (String.map (function ',' -> ' ' | c -> c) s)
  |> List.filter (fun w -> w <> "")

let parse_http_date value =
  let tokens =
    match words value with
    | first :: rest when String.length first > 0 && not (is_digit first.[0]) ->
        rest
    | tokens -> tokens
  in
  match tokens with
  | day :: month :: year :: time :: _ -> (
      match
        ( int_of_string_opt day,
          month_index month,
          int_of_string_opt year,
          String.split_on_char ':' time )
      with
      | Some day, Some month, Some year, [ h; mi; s ] -> (
          match
            (int_of_string_opt h, int_of_string_opt mi, int_of_string_opt s)
          with
          | Some h, Some mi, Some s ->
              Some
                ((float_of_int (days_from_civil year month day) *. 86400.)
                +. float_of_int ((h * 3600) + (mi * 60) + s))
          | _ -> None)
      | _ -> None)
  | _ -> None

let retry_after_seconds = function
  | None -> None
  | Some value -> (
      let value = String.trim value in
      match float_of_string_opt value with
      | Some seconds -> Some (Float.max 0. seconds)
      | None -> (
          match parse_http_date value with
          | Some target -> Some (Float.max 0. (target -. Unix.gettimeofday ()))
          | None -> None))

let header_value headers name =
  let name = String.lowercase_ascii name in
  List.fold_left
    (fun acc (k, v) ->
      if String.equal (String.lowercase_ascii (String.trim k)) name then
        Some (String.trim v)
      else acc)
    None headers

let retriable_status = [ 429; 500; 502; 503; 504 ]

let unreserved c =
  is_digit c
  || (c >= 'a' && c <= 'z')
  || (c >= 'A' && c <= 'Z')
  || c = '-' || c = '_' || c = '.' || c = '~'

let url_encode s =
  let buf = Buffer.create (String.length s) in
  String.iter
    (fun c ->
      if unreserved c then Buffer.add_char buf c
      else Buffer.add_string buf (Printf.sprintf "%%%02X" (Char.code c)))
    s;
  Buffer.contents buf

let strip_bom s =
  if String.length s >= 3 && String.sub s 0 3 = "\xef\xbb\xbf" then
    String.sub s 3 (String.length s - 3)
  else s

let csv_rows body =
  match Csv.input_all (Csv.of_string (strip_bom body)) with
  | [] -> []
  | header :: rows ->
      let width = List.length header in
      List.filter_map
        (fun row ->
          if List.for_all (fun cell -> String.trim cell = "") row then None
          else
            let padded =
              List.init width (fun i -> try List.nth row i with _ -> "")
            in
            Some (List.combine header padded))
        rows

let jitter = lazy (Random.State.make_self_init ())

module Client = struct
  type t = {
    base_url : string;
    timeout : float;
    max_attempts : int;
    backoff_base : float;
    backoff_max : float;
    total_timeout : float option;
    user_agent : string;
  }

  let rec strip_trailing_slash s =
    let n = String.length s in
    if n > 0 && s.[n - 1] = '/' then
      strip_trailing_slash (String.sub s 0 (n - 1))
    else s

  let create ?(base_url = default_url) ?(timeout = 10.) ?(max_attempts = 3)
      ?(backoff_base = 0.5) ?(backoff_max = 8.) ?total_timeout ?user_agent () =
    if max_attempts < 1 then
      invalid_arg "Cepis.Client.create: max_attempts must be >= 1";
    if timeout <= 0. then invalid_arg "Cepis.Client.create: timeout must be > 0";
    {
      base_url = strip_trailing_slash base_url;
      timeout;
      max_attempts;
      backoff_base;
      backoff_max;
      total_timeout;
      user_agent =
        (match user_agent with
        | Some ua -> ua
        | None -> Printf.sprintf "cepis-ocaml/%s (+https://cep.is)" version);
    }

  let base_url t = t.base_url
  let user_agent t = t.user_agent

  let perform t url =
    let config =
      Ezcurl.Config.(
        default |> follow_location true |> max_redirects 5 |> verbose false)
    in
    let set_opts curl =
      let ms = int_of_float (t.timeout *. 1000.) in
      Curl.set_timeoutms curl ms;
      Curl.set_connecttimeoutms curl ms
    in
    match
      Ezcurl.with_client ~set_opts (fun client ->
          Ezcurl.get ~client ~config ~tries:1
            ~headers:
              [
                ("User-Agent", t.user_agent);
                ("Accept", "application/json, text/csv;q=0.9, */*;q=0.1");
                ("Accept-Encoding", "identity");
              ]
            ~url ())
    with
    | Ok response ->
        Ok (response.Ezcurl.code, response.Ezcurl.body, response.Ezcurl.headers)
    | Error (_, message) -> Error message
    | exception Curl.CurlException (_, _, message) -> Error message

  let backoff t attempt retry_after =
    match retry_after with
    | Some seconds -> Float.min seconds t.backoff_max
    | None ->
        let ceiling =
          Float.min
            (t.backoff_base *. (2. ** float_of_int (attempt - 1)))
            t.backoff_max
        in
        if ceiling <= 0. then 0.
        else Random.State.float (Lazy.force jitter) ceiling

  let expired = function
    | None -> false
    | Some deadline -> Unix.gettimeofday () >= deadline

  let request t ?(accepted_status = [ 200 ]) path =
    let url = t.base_url ^ path in
    let deadline =
      match t.total_timeout with
      | None -> None
      | Some total -> Some (Unix.gettimeofday () +. total)
    in
    let attempts = t.max_attempts in
    let give_up status retry_after network =
      match status with
      | Some 429 ->
          Error
            (Rate_limited
               {
                 message =
                   Printf.sprintf
                     "rate limit reached at %s (after %d attempt(s))" url
                     attempts;
                 retry_after;
               })
      | Some status ->
          Error
            (Server_error
               {
                 message =
                   Printf.sprintf
                     "the API answered %d at %s (after %d attempt(s))" status
                     url attempts;
                 status;
               })
      | None ->
          Error
            (Network_error
               (Printf.sprintf "could not reach %s (after %d attempt(s)): %s"
                  url attempts
                  (match network with
                  | Some m -> m
                  | None -> "deadline reached")))
    in
    let wait_then t attempt retry_after =
      let wait = backoff t attempt retry_after in
      match deadline with
      | None ->
          if wait > 0. then Unix.sleepf wait;
          true
      | Some deadline ->
          let remaining = deadline -. Unix.gettimeofday () in
          if remaining <= 0. then false
          else begin
            let wait = Float.min wait remaining in
            if wait > 0. then Unix.sleepf wait;
            true
          end
    in
    let rec attempt n status retry_after network =
      if n > attempts || expired deadline then
        give_up status retry_after network
      else
        match perform t url with
        | Ok (code, body, _)
          when List.mem code accepted_status
               || not (List.mem code retriable_status) ->
            Ok (code, body)
        | Ok (code, _, headers) ->
            let retry_after =
              retry_after_seconds (header_value headers "Retry-After")
            in
            if n >= attempts then give_up (Some code) retry_after None
            else if wait_then t n retry_after then
              attempt (n + 1) (Some code) retry_after None
            else give_up (Some code) retry_after None
        | Error message ->
            if n >= attempts then give_up None None (Some message)
            else if wait_then t n None then
              attempt (n + 1) None None (Some message)
            else give_up None None (Some message)
    in
    attempt 1 None None None

  let ( let* ) = Result.bind

  let lookup t cep =
    let* number = normalize_cep cep in
    let* code, body = request t ("/" ^ number) in
    if code = 404 then
      Error
        (Cep_not_found
           {
             message =
               Printf.sprintf "CEP %s does not exist in the DNE database" number;
             faixa = faixa_of_body body;
           })
    else if code = 400 then
      Error
        (Invalid_cep
           (api_message body (Printf.sprintf "CEP %s rejected" number)))
    else
      let* fields = json_object_of_body body in
      Ok (address_of_fields fields)

  let lookup_chunk t ceps =
    let query = url_encode (String.concat "," ceps) in
    let* _, body = request t ("/csv?ceps=" ^ query) in
    let rows =
      List.map
        (fun row ->
          match List.assoc_opt "localidade" row with
          | Some place when String.trim place <> "" ->
              Some
                (address_of_fields
                   (List.map (fun (k, v) -> (k, `String v)) row))
          | _ -> None)
        (csv_rows body)
    in
    let wanted = List.length ceps in
    let got = List.length rows in
    if got = wanted then Ok rows
    else if got > wanted then Ok (List.filteri (fun i _ -> i < wanted) rows)
    else Ok (rows @ List.init (wanted - got) (fun _ -> None))

  let rec chunks size list =
    if list = [] then []
    else
      let head = List.filteri (fun i _ -> i < size) list in
      let tail = List.filteri (fun i _ -> i >= size) list in
      head :: chunks size tail

  let lookup_many t ceps =
    if ceps = [] then Ok []
    else
      List.fold_left
        (fun acc chunk ->
          let* acc = acc in
          let* rows = lookup_chunk t chunk in
          Ok (acc @ rows))
        (Ok []) (chunks max_batch ceps)

  let health t =
    let* _, body = request t ~accepted_status:[ 200; 503 ] "/health" in
    let* fields = json_object_of_body body in
    Ok (health_of_fields fields)
end

let default = lazy (Client.create ())
let lookup cep = Client.lookup (Lazy.force default) cep
let lookup_many ceps = Client.lookup_many (Lazy.force default) ceps
let health () = Client.health (Lazy.force default)