(** RFC 2617 HTTP Digest Authentication implementation *)

module Log = (val Logs.src_log (Logs.Src.create "requests.digest_auth" ~doc:"HTTP Digest Authentication") : Logs.LOG)

(** Digest auth challenge parameters from WWW-Authenticate header *)
type challenge = {
  realm : string;
  domain : string option;
  nonce : string;
  opaque : string option;
  stale : bool;
  algorithm : [`MD5 | `MD5_sess | `SHA256 | `SHA256_sess];
  qop : [`Auth | `Auth_int] list option;  (* quality of protection *)
  charset : string option;
  userhash : bool;
}

(** Client's chosen parameters for response *)
type client_data = {
  username : string;
  password : string;
  nc : int;  (* nonce count *)
  cnonce : string;  (* client nonce *)
  qop_chosen : [`Auth | `Auth_int] option;
}

(** Parse WWW-Authenticate header for Digest challenge *)
let parse_challenge header_value =
  (* Remove "Digest " prefix if present *)
  let value =
    if String.starts_with ~prefix:"Digest " header_value then
      String.sub header_value 7 (String.length header_value - 7)
    else header_value
  in

  (* Parse comma-separated key=value pairs *)
  let parse_params str =
    let rec parse_one pos acc =
      if pos >= String.length str then acc
      else
        (* Skip whitespace *)
        let pos = ref pos in
        while !pos < String.length str && str.[!pos] = ' ' do incr pos done;
        if !pos >= String.length str then acc
        else
          (* Find key *)
          let key_start = !pos in
          while !pos < String.length str && str.[!pos] <> '=' do incr pos done;
          if !pos >= String.length str then acc
          else
            let key = String.trim (String.sub str key_start (!pos - key_start)) in
            incr pos;  (* Skip '=' *)

            (* Parse value - may be quoted *)
            let value, next_pos =
              if !pos < String.length str && str.[!pos] = '"' then begin
                (* Quoted value *)
                incr pos;
                let value_start = !pos in
                while !pos < String.length str && str.[!pos] <> '"' do
                  if str.[!pos] = '\\' && !pos + 1 < String.length str then
                    pos := !pos + 2  (* Skip escaped character *)
                  else
                    incr pos
                done;
                let value = String.sub str value_start (!pos - value_start) in
                if !pos < String.length str then incr pos;  (* Skip closing quote *)
                (* Skip to next comma *)
                while !pos < String.length str && str.[!pos] <> ',' do incr pos done;
                if !pos < String.length str then incr pos;  (* Skip comma *)
                (value, !pos)
              end else begin
                (* Unquoted value *)
                let value_start = !pos in
                while !pos < String.length str && str.[!pos] <> ',' do incr pos done;
                let value = String.trim (String.sub str value_start (!pos - value_start)) in
                if !pos < String.length str then incr pos;  (* Skip comma *)
                (value, !pos)
              end
            in
            parse_one next_pos ((key, value) :: acc)
    in
    List.rev (parse_one 0 [])
  in

  let params = parse_params value in

  (* Extract required and optional parameters *)
  let get_param name = List.assoc_opt name params in
  let get_param_req name =
    match get_param name with
    | Some v -> v
    | None -> failwith (Printf.sprintf "Missing required Digest parameter: %s" name)
  in

  try
    let realm = get_param_req "realm" in
    let nonce = get_param_req "nonce" in

    let algorithm = match get_param "algorithm" with
      | Some "MD5" | None -> `MD5
      | Some "MD5-sess" -> `MD5_sess
      | Some "SHA-256" -> `SHA256
      | Some "SHA-256-sess" -> `SHA256_sess
      | Some a ->
          Log.warn (fun m -> m "Unknown digest algorithm: %s, using MD5" a);
          `MD5
    in

    let qop = match get_param "qop" with
      | None -> None
      | Some qop_str ->
          let qops = String.split_on_char ',' qop_str |> List.map String.trim in
          Some (List.filter_map (function
            | "auth" -> Some `Auth
            | "auth-int" -> Some `Auth_int
            | _ -> None
          ) qops)
    in

    Some {
      realm;
      domain = get_param "domain";
      nonce;
      opaque = get_param "opaque";
      stale = (match get_param "stale" with
        | Some "true" | Some "TRUE" -> true
        | _ -> false);
      algorithm;
      qop;
      charset = get_param "charset";
      userhash = (match get_param "userhash" with
        | Some "true" | Some "TRUE" -> true
        | _ -> false);
    }
  with
  | Failure msg ->
      Log.warn (fun m -> m "Failed to parse Digest challenge: %s" msg);
      None
  | Not_found -> None

(** Generate client nonce *)
let generate_cnonce () =
  let rand_bytes = Mirage_crypto_rng.generate 16 in
  Base64.encode_string rand_bytes

(** Hash function based on algorithm *)
let hash_function = function
  | `MD5 | `MD5_sess ->
      fun s -> Digestif.MD5.(to_hex (digest_string s))
  | `SHA256 | `SHA256_sess ->
      fun s -> Digestif.SHA256.(to_hex (digest_string s))

(** Calculate H(A1) according to RFC 2617 *)
let calculate_ha1 ~algorithm ~username ~realm ~password ~nonce ~cnonce =
  let hash = hash_function algorithm in
  match algorithm with
  | `MD5 | `SHA256 ->
      hash (Printf.sprintf "%s:%s:%s" username realm password)
  | `MD5_sess | `SHA256_sess ->
      let ha1_base = hash (Printf.sprintf "%s:%s:%s" username realm password) in
      hash (Printf.sprintf "%s:%s:%s" ha1_base nonce cnonce)

(** Calculate H(A2) according to RFC 2617 *)
let calculate_ha2 ~algorithm ~meth ~uri ~qop ~body =
  let hash = hash_function algorithm in
  let method_str = match meth with
    | `GET -> "GET" | `POST -> "POST" | `PUT -> "PUT"
    | `DELETE -> "DELETE" | `HEAD -> "HEAD" | `OPTIONS -> "OPTIONS"
    | `PATCH -> "PATCH" | `TRACE -> "TRACE" | `CONNECT -> "CONNECT"
    | `Other s -> s
  in
  match qop with
  | None | Some `Auth ->
      hash (Printf.sprintf "%s:%s" method_str (Uri.path_and_query uri))
  | Some `Auth_int ->
      (* For auth-int, include hash of entity body *)
      let body_hash = match body with
        | None -> hash ""
        | Some b -> hash b
      in
      hash (Printf.sprintf "%s:%s:%s" method_str (Uri.path_and_query uri) body_hash)

(** Calculate the response hash *)
let calculate_response ~ha1 ~ha2 ~nonce ~nc ~cnonce ~qop =
  let hash = hash_function `MD5 in  (* Response always uses the same hash as HA1 *)
  match qop with
  | None ->
      hash (Printf.sprintf "%s:%s:%s" ha1 nonce ha2)
  | Some qop_value ->
      let qop_str = match qop_value with
        | `Auth -> "auth"
        | `Auth_int -> "auth-int"
      in
      let nc_str = Printf.sprintf "%08x" nc in
      hash (Printf.sprintf "%s:%s:%s:%s:%s:%s" ha1 nonce nc_str cnonce qop_str ha2)

(** Generate Authorization header value for Digest auth *)
let generate_auth_header ~challenge ~client_data ~meth ~uri ~body =
  let { username; password; nc; cnonce; qop_chosen } = client_data in
  let { realm; nonce; opaque; algorithm; _ } = challenge in

  (* Calculate hashes *)
  let ha1 = calculate_ha1 ~algorithm ~username ~realm ~password ~nonce ~cnonce in
  let ha2 = calculate_ha2 ~algorithm ~meth ~uri ~qop:qop_chosen ~body in
  let response = calculate_response ~ha1 ~ha2 ~nonce ~nc ~cnonce ~qop:qop_chosen in

  (* Build Authorization header *)
  let params = [
    ("username", Printf.sprintf "\"%s\"" username);
    ("realm", Printf.sprintf "\"%s\"" realm);
    ("nonce", Printf.sprintf "\"%s\"" nonce);
    ("uri", Printf.sprintf "\"%s\"" (Uri.path_and_query uri));
    ("response", Printf.sprintf "\"%s\"" response);
  ] in

  let params = match algorithm with
    | `MD5 -> params  (* MD5 is default, don't need to specify *)
    | `MD5_sess -> ("algorithm", "MD5-sess") :: params
    | `SHA256 -> ("algorithm", "SHA-256") :: params
    | `SHA256_sess -> ("algorithm", "SHA-256-sess") :: params
  in

  let params = match opaque with
    | Some o -> ("opaque", Printf.sprintf "\"%s\"" o) :: params
    | None -> params
  in

  let params = match qop_chosen with
    | None -> params
    | Some qop ->
        let qop_str = match qop with `Auth -> "auth" | `Auth_int -> "auth-int" in
        let nc_str = Printf.sprintf "%08x" nc in
        ("qop", qop_str) ::
        ("nc", nc_str) ::
        ("cnonce", Printf.sprintf "\"%s\"" cnonce) ::
        params
  in

  "Digest " ^ String.concat ", " (List.map (fun (k, v) -> k ^ "=" ^ v) params)

(** Nonce counter storage - in production should be persistent *)
module NonceCounter = struct
  let table = Hashtbl.create 16

  let get_and_increment ~nonce =
    let current = try Hashtbl.find table nonce with Not_found -> 0 in
    Hashtbl.replace table nonce (current + 1);
    current + 1

  let reset ~nonce =
    Hashtbl.remove table nonce
end

(** Apply Digest authentication to a request *)
let apply_digest_auth ~username ~password ~meth ~uri ~headers ~body ~challenge_header =
  match parse_challenge challenge_header with
  | None ->
      Log.warn (fun m -> m "Failed to parse Digest challenge");
      headers
  | Some challenge ->
      (* Choose QOP if server offers options *)
      let qop_chosen = match challenge.qop with
        | None -> None
        | Some qops ->
            (* Prefer auth over auth-int for simplicity *)
            if List.mem `Auth qops then Some `Auth
            else if List.mem `Auth_int qops then Some `Auth_int
            else None
      in

      (* Get or generate client nonce *)
      let cnonce = generate_cnonce () in

      (* Get and increment nonce counter *)
      let nc = NonceCounter.get_and_increment ~nonce:challenge.nonce in

      let client_data = { username; password; nc; cnonce; qop_chosen } in
      let auth_value = generate_auth_header ~challenge ~client_data ~meth ~uri ~body in

      Cohttp.Header.add headers "Authorization" auth_value

(** Check if a response requires digest auth *)
let is_digest_challenge response =
  let status = Cohttp.Response.status response in
  match Cohttp.Code.code_of_status status with
  | 401 ->
      (match Cohttp.Header.get (Cohttp.Response.headers response) "www-authenticate" with
       | Some header when String.starts_with ~prefix:"Digest" header -> Some header
       | _ -> None)
  | _ -> None