let failures = ref 0
let total = ref 0

let check name ok =
  incr total;
  if ok then Printf.printf "ok   %s\n%!" name
  else begin
    incr failures;
    Printf.printf "FAIL %s\n%!" name
  end

type reply = { status : int; body : string; hdrs : (string * string) list }

let reply ?(hdrs = []) status body = { status; body; hdrs }

type server = {
  socket : Unix.file_descr;
  port : int;
  mutex : Mutex.t;
  mutable routes : (string * reply list) list;
  mutable handler : (string -> reply) option;
  mutable calls : string list;
}

let reason = function
  | 200 -> "OK"
  | 400 -> "Bad Request"
  | 404 -> "Not Found"
  | 429 -> "Too Many Requests"
  | 500 -> "Internal Server Error"
  | 503 -> "Service Unavailable"
  | _ -> "Status"

let contains_terminator s =
  let n = String.length s in
  let rec go i =
    if i + 4 > n then false
    else if String.sub s i 4 = "\r\n\r\n" then true
    else go (i + 1)
  in
  go 0

let read_request fd =
  let buf = Buffer.create 1024 in
  let chunk = Bytes.create 4096 in
  let rec go () =
    if contains_terminator (Buffer.contents buf) then Buffer.contents buf
    else
      let n = Unix.read fd chunk 0 4096 in
      if n = 0 then Buffer.contents buf
      else begin
        Buffer.add_subbytes buf chunk 0 n;
        go ()
      end
  in
  go ()

let request_path request =
  match String.index_opt request '\r' with
  | None -> "/"
  | Some i -> (
      match String.split_on_char ' ' (String.sub request 0 i) with
      | _ :: path :: _ -> path
      | _ -> "/")

let write_all fd s =
  let len = String.length s in
  let rec go off =
    if off < len then go (off + Unix.write_substring fd s off (len - off))
  in
  go 0

let next srv path =
  Mutex.lock srv.mutex;
  srv.calls <- path :: srv.calls;
  let answer =
    match srv.handler with
    | Some f -> f path
    | None -> (
        match List.assoc_opt path srv.routes with
        | None | Some [] -> reply 404 {|{"erro":"cep_nao_encontrado"}|}
        | Some [ only ] -> only
        | Some (head :: rest) ->
            srv.routes <- (path, rest) :: List.remove_assoc path srv.routes;
            head)
  in
  Mutex.unlock srv.mutex;
  answer

let serve srv fd =
  (try
     let path = request_path (read_request fd) in
     let answer = next srv path in
     let head = Buffer.create 256 in
     Buffer.add_string head
       (Printf.sprintf "HTTP/1.1 %d %s\r\n" answer.status (reason answer.status));
     List.iter
       (fun (k, v) -> Buffer.add_string head (Printf.sprintf "%s: %s\r\n" k v))
       answer.hdrs;
     Buffer.add_string head
       (Printf.sprintf "Content-Length: %d\r\n" (String.length answer.body));
     Buffer.add_string head "Content-Type: text/plain; charset=utf-8\r\n";
     Buffer.add_string head "Connection: close\r\n\r\n";
     write_all fd (Buffer.contents head ^ answer.body)
   with _ -> ());
  try Unix.close fd with _ -> ()

let start () =
  let socket = Unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
  Unix.setsockopt socket Unix.SO_REUSEADDR true;
  Unix.bind socket (Unix.ADDR_INET (Unix.inet_addr_loopback, 0));
  Unix.listen socket 32;
  let port =
    match Unix.getsockname socket with
    | Unix.ADDR_INET (_, port) -> port
    | _ -> 0
  in
  let srv =
    {
      socket;
      port;
      mutex = Mutex.create ();
      routes = [];
      handler = None;
      calls = [];
    }
  in
  let rec accept_loop () =
    match Unix.accept socket with
    | fd, _ ->
        ignore (Thread.create (fun () -> serve srv fd) ());
        accept_loop ()
    | exception _ -> ()
  in
  ignore (Thread.create accept_loop ());
  srv

let stop srv = try Unix.close srv.socket with _ -> ()
let url srv = Printf.sprintf "http://127.0.0.1:%d" srv.port
let respond srv path replies = srv.routes <- (path, replies) :: srv.routes

let count srv prefix =
  let n = String.length prefix in
  Mutex.lock srv.mutex;
  let matching =
    List.length
      (List.filter
         (fun c -> String.length c >= n && String.sub c 0 n = prefix)
         srv.calls)
  in
  Mutex.unlock srv.mutex;
  matching

let calls srv =
  Mutex.lock srv.mutex;
  let n = List.length srv.calls in
  Mutex.unlock srv.mutex;
  n

let with_server f =
  let srv = start () in
  let client =
    Cepis.Client.create ~base_url:(url srv) ~timeout:5.0 ~backoff_base:0.001
      ~backoff_max:0.01 ()
  in
  Fun.protect ~finally:(fun () -> stop srv) (fun () -> f srv client)

let se_fields =
  {|"cep":"01001000","logradouro":"Praça da Sé","complemento":"lado ímpar","unidade":"","bairro":"Sé","localidade":"São Paulo","uf":"SP","estado":"São Paulo","regiao":"Sudeste","ibge":"3550308","gia":"1004","ddd":"11","siafi":"7107"|}

let se_json = "{" ^ se_fields ^ "}"
let se_json_with extra = "{" ^ se_fields ^ "," ^ extra ^ "}"

let csv_header =
  "cep,logradouro,complemento,unidade,bairro,localidade,uf,estado,regiao,ibge,gia,ddd,siafi"

let csv rows = csv_header ^ "\n" ^ String.concat "\n" rows ^ "\n"

let is_error name predicate = function
  | Ok _ -> check name false
  | Error e -> check name (predicate e)

let contains haystack needle =
  let n = String.length haystack and m = String.length needle in
  let rec go i =
    if i + m > n then false
    else if String.sub haystack i m = needle then true
    else go (i + 1)
  in
  go 0

let test_normalisation () =
  List.iter
    (fun value ->
      check
        (Printf.sprintf "normalize_cep accepts %S" value)
        (Cepis.normalize_cep value = Ok "01001000"))
    [ "01001000"; "01001-000"; " 01001-000 "; "01001.000" ];
  List.iter
    (fun value ->
      check
        (Printf.sprintf "normalize_cep rejects %S" value)
        (match Cepis.normalize_cep value with
        | Error (Cepis.Invalid_cep _) -> true
        | _ -> false))
    [ ""; "123"; "0100100a"; "010010000"; "abcdefgh" ]

let test_successful_lookup () =
  with_server (fun srv client ->
      respond srv "/01001000" [ reply 200 se_json ];
      match Cepis.Client.lookup client "01001-000" with
      | Error _ -> check "successful lookup" false
      | Ok a ->
          check "lookup reads localidade" (a.Cepis.localidade = "São Paulo");
          check "lookup reads uf" (a.Cepis.uf = "SP");
          check "lookup masks the cep" (Cepis.formatted_cep a = "01001-000");
          check "lookup without geo yields None" (a.Cepis.geo = None))

let test_geo_present () =
  with_server (fun srv client ->
      respond srv "/01310100"
        [
          reply 200
            (se_json_with
               {|"geo":{"lat":-23.5,"lon":-46.6,"raio_m":120,"precisao":"logradouro"}|});
        ];
      match Cepis.Client.lookup client "01310100" with
      | Ok { Cepis.geo = Some g; _ } ->
          check "geo latitude is read" (Float.abs (g.Cepis.lat +. 23.5) < 1e-9);
          check "geo radius is read" (g.Cepis.raio_m = 120);
          check "street coordinate is precise" (Cepis.is_precise g)
      | _ -> check "geo present" false)

let test_geo_municipality () =
  with_server (fun srv client ->
      respond srv "/01001000"
        [
          reply 200
            (se_json_with
               {|"geo":{"lat":-23.5,"lon":-46.6,"raio_m":22000,"precisao":"localidade"}|});
        ];
      match Cepis.Client.lookup client "01001000" with
      | Ok { Cepis.geo = Some g; _ } ->
          check "municipality coordinate is not precise"
            (not (Cepis.is_precise g))
      | _ -> check "municipality coordinate" false)

let test_malformed_geo () =
  with_server (fun srv client ->
      respond srv "/01001000"
        [ reply 200 (se_json_with {|"geo":{"lat":"not a number"}|}) ];
      match Cepis.Client.lookup client "01001000" with
      | Ok a ->
          check "malformed geo is dropped" (a.Cepis.geo = None);
          check "malformed geo keeps the address"
            (a.Cepis.localidade = "São Paulo")
      | Error _ -> check "malformed geo" false)

let test_unknown_field () =
  with_server (fun srv client ->
      respond srv "/01001000"
        [ reply 200 (se_json_with {|"campo_do_futuro":"value"|}) ];
      match Cepis.Client.lookup client "01001000" with
      | Ok a ->
          check "unknown field is preserved in raw"
            (List.assoc_opt "campo_do_futuro" a.Cepis.raw
            = Some (`String "value"))
      | Error _ -> check "unknown field" false)

let test_not_found () =
  with_server (fun srv client ->
      respond srv "/99999999"
        [ reply 404 {|{"erro":"cep_nao_encontrado","mensagem":"gone"}|} ];
      let result = Cepis.Client.lookup client "99999999" in
      is_error "404 becomes Cep_not_found"
        (function Cepis.Cep_not_found _ -> true | _ -> false)
        result;
      is_error "a 404 without faixa yields None"
        (function
          | Cepis.Cep_not_found { faixa = None; _ } -> true | _ -> false)
        result;
      check "404 is not retried" (count srv "/99999999" = 1))

let areal_404 =
  {|{"erro":"cep_nao_encontrado","mensagem":"CEP não existe na base do DNE","faixa":{"localidade":"Areal","uf":"RJ","estado":"Rio de Janeiro","regiao":"Sudeste","ibge":"3300225","cep_ini":"25845001","cep_fim":"25849899","ceps":118}}|}

let test_not_found_with_faixa () =
  with_server (fun srv client ->
      respond srv "/25845000" [ reply 404 areal_404 ];
      match Cepis.Client.lookup client "25845000" with
      | Ok _ -> check "a faixa does not turn a 404 into an address" false
      | Error (Cepis.Cep_not_found { faixa = Some f; _ } as e) ->
          check "faixa names the municipality" (f.Cepis.localidade = "Areal");
          check "faixa carries the uf" (f.Cepis.uf = "RJ");
          check "faixa carries the ibge code" (f.Cepis.ibge = "3300225");
          check "faixa carries the region" (f.Cepis.regiao = "Sudeste");
          check "faixa bounds are the CEPs that exist, not the DNE's"
            (f.Cepis.cep_ini = "25845001" && f.Cepis.cep_fim = "25849899");
          check "faixa counts how sparse the range is" (f.Cepis.ceps = 118);
          check "the rendered error mentions the municipality"
            (contains (Cepis.string_of_error e) "Areal");
          check "the rendered error still says not found"
            (contains (Cepis.string_of_error e) "CEP not found")
      | Error _ -> check "a 404 with faixa is still Cep_not_found" false)

let test_not_found_faixa_is_defensive () =
  with_server (fun srv client ->
      List.iter
        (fun (path, body, label) ->
          respond srv path [ reply 404 body ];
          is_error
            (Printf.sprintf "%s yields no faixa, never a parse failure" label)
            (function
              | Cepis.Cep_not_found { faixa = None; _ } -> true | _ -> false)
            (Cepis.Client.lookup client (String.sub path 1 8)))
        [
          ( "/25845002",
            {|{"erro":"cep_nao_encontrado","faixa":null}|},
            "an explicit null" );
          ( "/25845003",
            {|{"erro":"cep_nao_encontrado","faixa":{}}|},
            "an empty object" );
          ( "/25845004",
            {|{"erro":"cep_nao_encontrado","faixa":"Areal"}|},
            "a bare string" );
          ( "/25845005",
            {|{"erro":"cep_nao_encontrado","faixa":{"uf":"RJ","ceps":"muitos"}}|},
            "a faixa with no municipality" );
          ("/25845006", "<html>captive portal</html>", "a body that is not JSON");
          ("/25845007", "", "an empty body");
        ])

let test_local_validation_makes_no_request () =
  with_server (fun srv client ->
      is_error "malformed CEP is rejected locally"
        (function Cepis.Invalid_cep _ -> true | _ -> false)
        (Cepis.Client.lookup client "123");
      check "malformed CEP never leaves the machine" (calls srv = 0))

let test_api_400 () =
  with_server (fun srv client ->
      respond srv "/00000000"
        [
          reply 400 {|{"erro":"cep_invalido","mensagem":"rejected by the API"}|};
        ];
      is_error "API 400 becomes Invalid_cep with the API message"
        (function
          | Cepis.Invalid_cep m -> m = "rejected by the API" | _ -> false)
        (Cepis.Client.lookup client "00000000");
      check "400 is not retried" (count srv "/00000000" = 1))

let test_not_json () =
  with_server (fun srv client ->
      respond srv "/01001000" [ reply 200 "<html>captive portal</html>" ];
      is_error "unparseable body becomes Invalid_response"
        (function Cepis.Invalid_response _ -> true | _ -> false)
        (Cepis.Client.lookup client "01001000"))

let test_retry_then_success () =
  with_server (fun srv client ->
      respond srv "/01001000"
        [ reply 500 {|{"erro":"interno"}|}; reply 200 se_json ];
      (match Cepis.Client.lookup client "01001000" with
      | Ok a -> check "retry then success" (a.Cepis.localidade = "São Paulo")
      | Error _ -> check "retry then success" false);
      check "500 was retried once" (count srv "/01001000" = 2))

let test_retry_exhausted () =
  with_server (fun srv client ->
      respond srv "/01001000" [ reply 503 {|{"erro":"interno"}|} ];
      is_error "exhausted attempts become Server_error"
        (function
          | Cepis.Server_error { status; _ } -> status = 503 | _ -> false)
        (Cepis.Client.lookup client "01001000");
      check "three attempts is the default" (count srv "/01001000" = 3))

let test_max_attempts_one () =
  let srv = start () in
  Fun.protect
    ~finally:(fun () -> stop srv)
    (fun () ->
      let client =
        Cepis.Client.create ~base_url:(url srv) ~max_attempts:1 ~timeout:5.0 ()
      in
      respond srv "/01001000" [ reply 500 "{}" ];
      is_error "max_attempts:1 still reports the error"
        (function Cepis.Server_error _ -> true | _ -> false)
        (Cepis.Client.lookup client "01001000");
      check "max_attempts:1 disables retrying" (count srv "/01001000" = 1))

let test_rate_limited () =
  with_server (fun srv client ->
      respond srv "/01001000" [ reply ~hdrs:[ ("Retry-After", "7") ] 429 "{}" ];
      is_error "429 becomes Rate_limited carrying retry_after"
        (function
          | Cepis.Rate_limited { retry_after = Some s; _ } ->
              Float.abs (s -. 7.) < 1e-9
          | _ -> false)
        (Cepis.Client.lookup client "01001000"))

let test_retry_after_http_date () =
  with_server (fun srv client ->
      respond srv "/01001000"
        [
          reply
            ~hdrs:[ ("Retry-After", "Wed, 21 Oct 2026 07:28:00 GMT") ]
            429 "{}";
        ];
      is_error "Retry-After as an HTTP date is understood"
        (function
          | Cepis.Rate_limited { retry_after = Some _; _ } -> true | _ -> false)
        (Cepis.Client.lookup client "01001000"))

let test_retry_after_past_date () =
  with_server (fun srv client ->
      respond srv "/01001000"
        [
          reply
            ~hdrs:[ ("Retry-After", "Wed, 21 Oct 1998 07:28:00 GMT") ]
            429 "{}";
        ];
      is_error "an HTTP date already past asks for no wait"
        (function
          | Cepis.Rate_limited { retry_after = Some s; _ } -> s = 0.
          | _ -> false)
        (Cepis.Client.lookup client "01001000"))

let test_retry_after_cap () =
  with_server (fun srv client ->
      respond srv "/01001000"
        [ reply ~hdrs:[ ("Retry-After", "3600") ] 429 "{}" ];
      let started = Unix.gettimeofday () in
      ignore (Cepis.Client.lookup client "01001000");
      check "an absurd Retry-After is capped"
        (Unix.gettimeofday () -. started < 5.))

let test_network_error () =
  let client =
    Cepis.Client.create ~base_url:"http://127.0.0.1:1" ~max_attempts:2
      ~backoff_base:0.001 ~timeout:1.0 ()
  in
  is_error "a refused connection becomes Network_error"
    (function Cepis.Network_error _ -> true | _ -> false)
    (Cepis.Client.lookup client "01001000")

let test_total_timeout () =
  let srv = start () in
  Fun.protect
    ~finally:(fun () -> stop srv)
    (fun () ->
      let client =
        Cepis.Client.create ~base_url:(url srv) ~max_attempts:50
          ~backoff_base:0.05 ~backoff_max:0.05 ~total_timeout:0.3 ()
      in
      respond srv "/01001000" [ reply 500 "{}" ];
      let started = Unix.gettimeofday () in
      is_error "total_timeout still reports the last failure"
        (function Cepis.Server_error _ -> true | _ -> false)
        (Cepis.Client.lookup client "01001000");
      check "total_timeout bounds the whole call"
        (Unix.gettimeofday () -. started < 3.);
      check "total_timeout cuts the attempts short" (count srv "/01001000" < 50))

let test_batch_order () =
  with_server (fun srv client ->
      respond srv "/csv?ceps=01001000%2C99999999%2C20040020"
        [
          reply 200
            (csv
               [
                 "01001000,Praça da Sé,lado ímpar,,Sé,São Paulo,SP,São \
                  Paulo,Sudeste,3550308,1004,11,7107";
                 "99999999,,,,,,,,,,,,";
                 "20040020,Avenida Rio Branco,,,Centro,Rio de Janeiro,RJ,Rio \
                  de Janeiro,Sudeste,3304557,,21,6001";
               ]);
        ];
      match
        Cepis.Client.lookup_many client [ "01001000"; "99999999"; "20040020" ]
      with
      | Error _ -> check "batch order" false
      | Ok rows -> (
          check "batch keeps the input length" (List.length rows = 3);
          (match List.nth rows 0 with
          | Some a ->
              check "batch resolves the first" (a.Cepis.localidade = "São Paulo")
          | None -> check "batch resolves the first" false);
          check "a missing CEP is None at its position" (List.nth rows 1 = None);
          match List.nth rows 2 with
          | Some a -> check "batch resolves the third" (a.Cepis.uf = "RJ")
          | None -> check "batch resolves the third" false))

let test_batch_no_geo () =
  with_server (fun srv client ->
      respond srv "/csv?ceps=01001000"
        [
          reply 200
            (csv
               [
                 "01001000,Praça da Sé,,,Sé,São Paulo,SP,São \
                  Paulo,Sudeste,3550308,1004,11,7107";
               ]);
        ];
      match Cepis.Client.lookup_many client [ "01001000" ] with
      | Ok [ Some a ] -> check "batches carry no geo" (a.Cepis.geo = None)
      | _ -> check "batches carry no geo" false)

let test_batch_quoted_comma () =
  with_server (fun srv client ->
      respond srv "/csv?ceps=01001000"
        [
          reply 200
            (csv
               [
                 {|01001000,"Rua Um, Dois",,,Centro,São Paulo,SP,São Paulo,Sudeste,1,,11,1|};
               ]);
        ];
      match Cepis.Client.lookup_many client [ "01001000" ] with
      | Ok [ Some a ] ->
          check "a quoted comma is parsed" (a.Cepis.logradouro = "Rua Um, Dois")
      | _ -> check "a quoted comma is parsed" false)

let test_batch_empty () =
  with_server (fun srv client ->
      check "an empty batch returns an empty list"
        (Cepis.Client.lookup_many client [] = Ok []);
      check "an empty batch makes no request" (calls srv = 0))

let test_batch_chunking () =
  with_server (fun srv client ->
      let ceps = List.init 501 (fun i -> Printf.sprintf "%08d" i) in
      let row c = Printf.sprintf "%s,,,,,X,SP,São Paulo,Sudeste,1,,11,1" c in
      let requests = ref 0 in
      srv.handler <-
        Some
          (fun _ ->
            incr requests;
            let rows =
              if !requests = 1 then
                List.filteri (fun i _ -> i < 500) ceps |> List.map row
              else List.filteri (fun i _ -> i >= 500) ceps |> List.map row
            in
            reply 200 (csv rows));
      match Cepis.Client.lookup_many client ceps with
      | Error _ -> check "batch chunking" false
      | Ok rows ->
          check "501 CEPs become two requests" (!requests = 2);
          check "chunked results keep the input length" (List.length rows = 501))

let test_batch_truncated () =
  with_server (fun srv client ->
      respond srv "/csv?ceps=01001000%2C20040020"
        [
          reply 200
            (csv [ "01001000,,,,,São Paulo,SP,São Paulo,Sudeste,1,,11,1" ]);
        ];
      match Cepis.Client.lookup_many client [ "01001000"; "20040020" ] with
      | Ok rows ->
          check "a short response is padded, not shifted"
            (List.length rows = 2 && List.nth rows 1 = None)
      | Error _ -> check "a short response is padded, not shifted" false)

let test_health () =
  with_server (fun srv client ->
      respond srv "/health"
        [
          reply 200
            {|{"status":"ok","dne":"26071","ceps":1596168,"geo":894489,"base_gerada_em":"2026-07-24T03:48:28Z"}|};
        ];
      match Cepis.Client.health client with
      | Ok h ->
          check "health reports ok" (Cepis.is_ok h);
          check "health reports the database" (h.Cepis.dne = "26071");
          check "health counts the CEPs" (h.Cepis.ceps = 1596168)
      | Error _ -> check "health" false)

let test_health_degraded () =
  with_server (fun srv client ->
      respond srv "/health"
        [ reply 503 {|{"status":"degradado","dne":"","ceps":0,"geo":0}|} ];
      match Cepis.Client.health client with
      | Ok h ->
          check "a degraded API is reported, not an error" (not (Cepis.is_ok h))
      | Error _ -> check "a degraded API is reported, not an error" false)

let test_configuration () =
  check "a trailing slash is stripped"
    (Cepis.Client.base_url (Cepis.Client.create ~base_url:"http://x/" ())
    = "http://x");
  check "the user agent identifies the library"
    (let ua = Cepis.Client.user_agent (Cepis.Client.create ()) in
     String.length ua > 12 && String.sub ua 0 12 = "cepis-ocaml/");
  check "max_attempts below 1 is refused"
    (match Cepis.Client.create ~max_attempts:0 () with
    | _ -> false
    | exception Invalid_argument _ -> true);
  check "a non-positive timeout is refused"
    (match Cepis.Client.create ~timeout:0. () with
    | _ -> false
    | exception Invalid_argument _ -> true)

let () =
  test_normalisation ();
  test_successful_lookup ();
  test_geo_present ();
  test_geo_municipality ();
  test_malformed_geo ();
  test_unknown_field ();
  test_not_found ();
  test_not_found_with_faixa ();
  test_not_found_faixa_is_defensive ();
  test_local_validation_makes_no_request ();
  test_api_400 ();
  test_not_json ();
  test_retry_then_success ();
  test_retry_exhausted ();
  test_max_attempts_one ();
  test_rate_limited ();
  test_retry_after_http_date ();
  test_retry_after_past_date ();
  test_retry_after_cap ();
  test_network_error ();
  test_total_timeout ();
  test_batch_order ();
  test_batch_no_geo ();
  test_batch_quoted_comma ();
  test_batch_empty ();
  test_batch_chunking ();
  test_batch_truncated ();
  test_health ();
  test_health_degraded ();
  test_configuration ();
  Printf.printf "\n%d checks, %d failure(s)\n%!" !total !failures;
  if !failures > 0 then exit 1