open Fastcgi_record

(** {1 Request Roles} *)

type role =
  | Responder
  | Authorizer
  | Filter

let pp_role ppf = function
  | Responder -> Format.pp_print_string ppf "Responder"
  | Authorizer -> Format.pp_print_string ppf "Authorizer"
  | Filter -> Format.pp_print_string ppf "Filter"

let role_of_begin_request record =
  if record.record_type <> Begin_request then
    Error "Expected BEGIN_REQUEST record"
  else if String.length record.content <> 8 then
    Error "Invalid BEGIN_REQUEST content length"
  else
    let content = record.content in
    let role_int = (Char.code content.[0] lsl 8) lor (Char.code content.[1]) in
    match role_int with
    | 1 -> Ok Responder
    | 2 -> Ok Authorizer
    | 3 -> Ok Filter
    | n -> Error (Printf.sprintf "Unknown FastCGI role: %d" n)

(** {1 Request Context} *)

type t = {
  request_id : request_id;
  role : role;
  keep_conn : bool;
  params : KV.t;
  stdin_data : string;
  data_stream : string option;
}

let pp ppf request =
  let data_str = match request.data_stream with
    | None -> "None"
    | Some d -> Printf.sprintf "Some(%d bytes)" (String.length d)
  in
  Format.fprintf ppf
    "@[<2>{ request_id = %d;@ role = %a;@ keep_conn = %b;@ params = %a;@ stdin = %d bytes;@ data = %s }@]"
    request.request_id
    pp_role request.role
    request.keep_conn
    (KV.pp) request.params
    (String.length request.stdin_data)
    data_str

let create record =
  match role_of_begin_request record with
  | Error _ as e -> e
  | Ok role ->
      if String.length record.content <> 8 then
        Error "Invalid BEGIN_REQUEST content length"
      else
        let flags_int = Char.code record.content.[2] in
        let keep_conn = (flags_int land 1) <> 0 in
        Ok {
          request_id = record.request_id;
          role;
          keep_conn;
          params = KV.empty;
          stdin_data = "";
          data_stream = if role = Filter then Some "" else None;
        }


(** {1 Stream Processing} *)

(** Helper functions for result binding to simplify nested pattern matching *)
let ( let* ) = Result.bind

let is_stream_terminator record =
  String.length record.content = 0


let read_params buf_read =
  Printf.eprintf "[DEBUG] read_params_from_flow: Starting\n%!";
  let params = ref KV.empty in
  let rec loop () =
    try
      Printf.eprintf "[DEBUG] read_params_from_flow: Reading next PARAMS record\n%!";
      let record = Fastcgi_record.read buf_read in
      Printf.eprintf "[DEBUG] read_params_from_flow: Got record type=%s, content_length=%d\n%!"
        (Format.asprintf "%a" pp_record record.record_type)
        (String.length record.content);
      if record.record_type <> Params then
        Error (Printf.sprintf "Expected PARAMS record, got %s"
                      (Format.asprintf "%a" pp_record record.record_type))
      else if is_stream_terminator record then (
        Printf.eprintf "[DEBUG] read_params_from_flow: Got stream terminator, returning %d params\n%!"
          (Fastcgi_record.KV.cardinal !params);
        Ok !params
      ) else (
        let record_params = KV.decode record.content in
        Printf.eprintf "[DEBUG] read_params_from_flow: Decoded %d params from record\n%!"
          (Fastcgi_record.KV.cardinal record_params);
        params := KV.to_seq record_params
                  |> Seq.fold_left (fun acc (k, v) -> KV.add k v acc) !params;
        loop ()
      )
    with
    | End_of_file -> 
        Printf.eprintf "[DEBUG] read_params_from_flow: Hit End_of_file\n%!";
        Error "Unexpected end of stream while reading PARAMS"
    | exn -> 
        Printf.eprintf "[DEBUG] read_params_from_flow: Exception: %s\n%!" (Printexc.to_string exn);
        Error (Printf.sprintf "Error reading PARAMS: %s" (Printexc.to_string exn))
  in
  loop ()

let read_stdin buf_read =
  Printf.eprintf "[DEBUG] read_stdin_from_flow: Starting\n%!";
  let data = Buffer.create 1024 in
  let rec loop () =
    try
      Printf.eprintf "[DEBUG] read_stdin_from_flow: Reading next STDIN record\n%!";
      let record = Fastcgi_record.read buf_read in
      Printf.eprintf "[DEBUG] read_stdin_from_flow: Got record type=%s, content_length=%d\n%!"
        (Format.asprintf "%a" pp_record record.record_type)
        (String.length record.content);
      if record.record_type <> Stdin then
        Error (Printf.sprintf "Expected STDIN record, got %s"
                      (Format.asprintf "%a" pp_record record.record_type))
      else if is_stream_terminator record then (
        Printf.eprintf "[DEBUG] read_stdin_from_flow: Got stream terminator, total stdin=%d bytes\n%!"
          (Buffer.length data);
        Ok (Buffer.contents data)
      ) else (
        Buffer.add_string data record.content;
        Printf.eprintf "[DEBUG] read_stdin_from_flow: Added %d bytes, total now %d\n%!"
          (String.length record.content) (Buffer.length data);
        loop ()
      )
    with
    | End_of_file -> 
        Printf.eprintf "[DEBUG] read_stdin_from_flow: Hit End_of_file\n%!";
        Error "Unexpected end of stream while reading STDIN"
    | exn -> 
        Printf.eprintf "[DEBUG] read_stdin_from_flow: Exception: %s\n%!" (Printexc.to_string exn);
        Error (Printf.sprintf "Error reading STDIN: %s" (Printexc.to_string exn))
  in
  loop ()

(** Read DATA stream for Filter role *)
let read_data buf_read =
  let data_buf = Buffer.create 1024 in
  let rec read_data () =
    try
      let record = Fastcgi_record.read buf_read in
      if record.record_type <> Data then
        Error "Expected DATA record"
      else if is_stream_terminator record then
        Ok (Buffer.contents data_buf)
      else (
        Buffer.add_string data_buf record.content;
        read_data ()
      )
    with
    | End_of_file -> Error "Unexpected end of DATA stream"
    | exn -> Error (Printf.sprintf "Error reading DATA: %s" (Printexc.to_string exn))
  in
  read_data ()

(** Read request streams based on role *)
let read_request_streams request buf_read =
  Printf.eprintf "[DEBUG] read_request_streams: Processing role=%s\n%!" 
    (Format.asprintf "%a" pp_role request.role);
  match request.role with
  | Authorizer -> 
      Printf.eprintf "[DEBUG] read_request_streams: Authorizer role, no streams to read\n%!";
      Ok request
  | Responder ->
      Printf.eprintf "[DEBUG] read_request_streams: Responder role, reading STDIN\n%!";
      let* stdin_data = read_stdin buf_read in
      Printf.eprintf "[DEBUG] read_request_streams: Got STDIN data, %d bytes\n%!" (String.length stdin_data);
      Ok { request with stdin_data }
  | Filter ->
      Printf.eprintf "[DEBUG] read_request_streams: Filter role, reading STDIN and DATA\n%!";
      let* stdin_data = read_stdin buf_read in
      Printf.eprintf "[DEBUG] read_request_streams: Got STDIN data, %d bytes\n%!" (String.length stdin_data);
      let request = { request with stdin_data } in
      let* data = read_data buf_read in
      Printf.eprintf "[DEBUG] read_request_streams: Got DATA stream, %d bytes\n%!" (String.length data);
      Ok { request with data_stream = Some data }

let read_request buf_read =
  Printf.eprintf "[DEBUG] read_request: Starting\n%!";
  try
    (* Read BEGIN_REQUEST *)
    Printf.eprintf "[DEBUG] read_request: Reading BEGIN_REQUEST record\n%!";
    let begin_record = Fastcgi_record.read buf_read in
    Printf.eprintf "[DEBUG] read_request: Got BEGIN_REQUEST record: %s\n%!"
      (Format.asprintf "%a" (Fastcgi_record.pp ~max_content_len:50) begin_record);
    let* request = create begin_record in
    Printf.eprintf "[DEBUG] read_request: Created request with role=%s, id=%d\n%!"
      (Format.asprintf "%a" pp_role request.role) request.request_id;
    (* Read PARAMS stream *)
    Printf.eprintf "[DEBUG] read_request: Reading PARAMS stream\n%!";
    let* params = read_params buf_read in
    Printf.eprintf "[DEBUG] read_request: Got %d params\n%!" (Fastcgi_record.KV.cardinal params);
    let request = { request with params } in
    (* Read remaining streams based on role *)
    Printf.eprintf "[DEBUG] read_request: Reading streams for role=%s\n%!"
      (Format.asprintf "%a" pp_role request.role);
    let result = read_request_streams request buf_read in
    Printf.eprintf "[DEBUG] read_request: Finished reading request\n%!";
    result
  with
  | End_of_file -> 
      Printf.eprintf "[DEBUG] read_request: Hit End_of_file\n%!";
      Error "Unexpected end of stream"
  | exn -> 
      Printf.eprintf "[DEBUG] read_request: Exception: %s\n%!" (Printexc.to_string exn);
      Error (Printf.sprintf "Error reading request: %s" (Printexc.to_string exn))

(** {1 Response Generation} *)

type app_status = int
type protocol_status = 
  | Request_complete
  | Cant_mpx_conn  
  | Overloaded
  | Unknown_role

let pp_protocol_status ppf = function
  | Request_complete -> Format.pp_print_string ppf "Request_complete"
  | Cant_mpx_conn -> Format.pp_print_string ppf "Cant_mpx_conn"
  | Overloaded -> Format.pp_print_string ppf "Overloaded"
  | Unknown_role -> Format.pp_print_string ppf "Unknown_role"

let protocol_status_to_int = function
  | Request_complete -> 0
  | Cant_mpx_conn -> 1
  | Overloaded -> 2
  | Unknown_role -> 3

let write_stream_records buf_write request_id record_type content =
  let max_chunk = 65535 in (* FastCGI max record content length *)
  let len = String.length content in
  let rec chunk_string pos =
    if pos < len then begin
      let chunk_len = min max_chunk (len - pos) in
      let record = Fastcgi_record.create ~record:record_type ~request_id ~content ~offset:pos ~length:chunk_len () in
      Fastcgi_record.write buf_write record;
      chunk_string (pos + chunk_len)
    end
  in
  chunk_string 0;
  let terminator = Fastcgi_record.create ~record:record_type ~request_id ~content:"" () in
  Fastcgi_record.write buf_write terminator

let write_stdout_records buf_write request_id content =
  Printf.eprintf "[DEBUG] write_stdout_records: Writing %d bytes for request_id=%d\n%!" 
    (String.length content) request_id;
  write_stream_records buf_write request_id Stdout content

let write_stderr_records buf_write request_id content =
  Printf.eprintf "[DEBUG] write_stderr_records: Writing %d bytes for request_id=%d\n%!" 
    (String.length content) request_id;
  write_stream_records buf_write request_id Stderr content

let write_end_request buf_write request_id app_status protocol_status =
  let content = 
    let buf = Bytes.create 8 in
    Bytes.set_int32_be buf 0 (Int32.of_int app_status);
    Bytes.set_uint8 buf 4 (protocol_status_to_int protocol_status);
    Bytes.set_uint8 buf 5 0; (* reserved *)
    Bytes.set_uint8 buf 6 0; (* reserved *)
    Bytes.set_uint8 buf 7 0; (* reserved *)
    Bytes.to_string buf
  in
  let record = Fastcgi_record.create ~record:End_request ~request_id ~content () in
  Fastcgi_record.write buf_write record