(*---------------------------------------------------------------------------
  Copyright (c) 2025 Anil Madhavapeddy <anil@recoil.org>. All rights reserved.
  SPDX-License-Identifier: ISC
 ---------------------------------------------------------------------------*)

let src = Logs.Src.create "requests.headers" ~doc:"HTTP Headers"
module Log = (val Logs.src_log src : Logs.LOG)

(* Use a map with lowercase keys for case-insensitive lookup *)
module StringMap = Map.Make(String)

type t = (string * string list) StringMap.t

let empty = StringMap.empty

let normalize_key k = String.lowercase_ascii k

(** {1 Header Injection Prevention}

    Per Recommendation #3: Validate that header names and values do not contain
    newlines (CR/LF) which could enable HTTP request smuggling attacks.

    Note: We use Invalid_argument here to avoid a dependency cycle with Error module.
    The error will be caught and wrapped appropriately by higher-level code. *)

exception Invalid_header of { name: string; reason: string }

let validate_header_name name =
  if String.contains name '\r' || String.contains name '\n' then
    raise (Invalid_header {
      name;
      reason = "Header name contains CR/LF characters (potential HTTP smuggling)"
    })

let validate_header_value name value =
  if String.contains value '\r' || String.contains value '\n' then
    raise (Invalid_header {
      name;
      reason = "Header value contains CR/LF characters (potential HTTP smuggling)"
    })

let add key value t =
  validate_header_name key;
  validate_header_value key value;
  let nkey = normalize_key key in
  let existing =
    match StringMap.find_opt nkey t with
    | Some (_, values) -> values
    | None -> []
  in
  (* Append to maintain order, avoiding reversal on retrieval *)
  StringMap.add nkey (key, existing @ [value]) t

let set key value t =
  validate_header_name key;
  validate_header_value key value;
  let nkey = normalize_key key in
  StringMap.add nkey (key, [value]) t

let get key t =
  let nkey = normalize_key key in
  match StringMap.find_opt nkey t with
  | Some (_, values) -> List.nth_opt values 0
  | None -> None

let get_all key t =
  let nkey = normalize_key key in
  match StringMap.find_opt nkey t with
  | Some (_, values) -> values
  | None -> []

let remove key t =
  let nkey = normalize_key key in
  StringMap.remove nkey t

let mem key t =
  let nkey = normalize_key key in
  StringMap.mem nkey t

let of_list lst =
  List.fold_left (fun acc (k, v) -> add k v acc) empty lst

let to_list t =
  StringMap.fold (fun _ (orig_key, values) acc ->
    (* Values are already in correct order, build list in reverse then reverse at end *)
    List.fold_left (fun acc v -> (orig_key, v) :: acc) acc values
  ) t []
  |> List.rev

let merge t1 t2 =
  StringMap.union (fun _ _ v2 -> Some v2) t1 t2

(* Common header builders *)

let content_type mime t =
  set "Content-Type" (Mime.to_string mime) t

let content_length len t =
  set "Content-Length" (Int64.to_string len) t

let accept mime t =
  set "Accept" (Mime.to_string mime) t

let authorization value t =
  set "Authorization" value t

let bearer token t =
  set "Authorization" (Printf.sprintf "Bearer %s" token) t

let basic ~username ~password t =
  let credentials = Printf.sprintf "%s:%s" username password in
  let encoded = Base64.encode_exn credentials in
  set "Authorization" (Printf.sprintf "Basic %s" encoded) t

let user_agent ua t =
  set "User-Agent" ua t

let host h t =
  set "Host" h t

let cookie name value t =
  add "Cookie" (Printf.sprintf "%s=%s" name value) t

let range ~start ?end_ () t =
  let range_value = match end_ with
    | None -> Printf.sprintf "bytes=%Ld-" start
    | Some e -> Printf.sprintf "bytes=%Ld-%Ld" start e
  in
  set "Range" range_value t

(** {1 HTTP 100-Continue Support}

    Per Recommendation #7: Expect: 100-continue protocol for large uploads.
    RFC 9110 Section 10.1.1 (Expect) *)

let expect value t =
  set "Expect" value t

let expect_100_continue t =
  set "Expect" "100-continue" t

(** {1 Cache Control Headers}

    Per Recommendation #17 and #19: Response caching and conditional requests.
    RFC 9111 (HTTP Caching), RFC 9110 Section 8.8.2-8.8.3 (Last-Modified, ETag) *)

let if_none_match etag t =
  set "If-None-Match" etag t

let if_match etag t =
  set "If-Match" etag t

let if_modified_since date t =
  set "If-Modified-Since" date t

let if_unmodified_since date t =
  set "If-Unmodified-Since" date t

(** Format a Ptime.t as an HTTP-date (RFC 9110 Section 5.6.7) *)
let http_date_of_ptime time =
  (* HTTP-date format: "Sun, 06 Nov 1994 08:49:37 GMT" *)
  let (year, month, day), ((hour, min, sec), _tz_offset) = Ptime.to_date_time time in
  let weekday = match Ptime.weekday time with
    | `Sun -> "Sun" | `Mon -> "Mon" | `Tue -> "Tue" | `Wed -> "Wed"
    | `Thu -> "Thu" | `Fri -> "Fri" | `Sat -> "Sat"
  in
  let month_name = [| ""; "Jan"; "Feb"; "Mar"; "Apr"; "May"; "Jun";
                      "Jul"; "Aug"; "Sep"; "Oct"; "Nov"; "Dec" |].(month) in
  Printf.sprintf "%s, %02d %s %04d %02d:%02d:%02d GMT"
    weekday day month_name year hour min sec

let if_modified_since_ptime time t =
  if_modified_since (http_date_of_ptime time) t

let if_unmodified_since_ptime time t =
  if_unmodified_since (http_date_of_ptime time) t

let cache_control directives t =
  set "Cache-Control" directives t

(** Build Cache-Control header from common directive components.
    For max_stale: [None] = not present, [Some None] = any staleness, [Some (Some n)] = n seconds *)
let cache_control_directives
    : ?max_age:int ->
      ?max_stale:int option option ->
      ?min_fresh:int ->
      ?no_cache:bool ->
      ?no_store:bool ->
      ?no_transform:bool ->
      ?only_if_cached:bool ->
      unit -> t -> t
  = fun
    ?max_age
    ?max_stale
    ?min_fresh
    ?(no_cache = false)
    ?(no_store = false)
    ?(no_transform = false)
    ?(only_if_cached = false)
    () t ->
  let directives = [] in
  let directives = match max_age with
    | Some age -> Printf.sprintf "max-age=%d" age :: directives
    | None -> directives
  in
  let directives = match max_stale with
    | Some (Some None) -> "max-stale" :: directives
    | Some (Some (Some secs)) -> Printf.sprintf "max-stale=%d" secs :: directives
    | Some None | None -> directives
  in
  let directives = match min_fresh with
    | Some secs -> Printf.sprintf "min-fresh=%d" secs :: directives
    | None -> directives
  in
  let directives = if no_cache then "no-cache" :: directives else directives in
  let directives = if no_store then "no-store" :: directives else directives in
  let directives = if no_transform then "no-transform" :: directives else directives in
  let directives = if only_if_cached then "only-if-cached" :: directives else directives in
  match directives with
  | [] -> t
  | _ -> set "Cache-Control" (String.concat ", " (List.rev directives)) t

let etag value t =
  set "ETag" value t

let last_modified date t =
  set "Last-Modified" date t

let last_modified_ptime time t =
  last_modified (http_date_of_ptime time) t

(* Additional helper for getting multiple header values *)
let get_multi key t = get_all key t

(* Pretty printer for headers *)
let pp ppf t =
  Format.fprintf ppf "@[<v>Headers:@,";
  let headers = to_list t in
  List.iter (fun (k, v) ->
    Format.fprintf ppf "  %s: %s@," k v
  ) headers;
  Format.fprintf ppf "@]"

let pp_brief ppf t =
  let headers = to_list t in
  let count = List.length headers in
  Format.fprintf ppf "Headers(%d entries)" count