(* server.ml - Eio static file server using httpzo *)

open Eio.Std

(* Response buffer size - 64KB for headers *)
let response_buffer_size = 65536

type config = {
  port : int;
  root : Eio.Fs.dir_ty Eio.Path.t;
  limits : Buf_read.limits;
}

let default_config ~fs = {
  port = 8080;
  root = fs;
  limits = Buf_read.default_limits;
}

(* Connection state - generic over flow type for accept_fork compatibility *)
type 'a conn_state =
  { flow : 'a
  ; read_buf : Base_bigstring.t
  ; write_buf : Base_bigstring.t
  ; mutable read_len : int
  ; mutable keep_alive : bool
  (* Reusable arrays for range parsing *)
  ; ranges : Range.byte_range array
  ; resolved : Range.resolved array
  }

(* Create connection state *)
let create_conn flow =
  { flow
  ; read_buf = Buf_read.create ()
  ; write_buf =
      Bigarray.Array1.create Bigarray.char Bigarray.c_layout response_buffer_size
  ; read_len = 0
  ; keep_alive = true
  ; ranges = Array.make Range.max_ranges Range.empty
  ; resolved = Array.make Range.max_ranges Range.empty_resolved
  }

(* MIME type detection using magic-mime *)
let mime_type_of_path path = Magic_mime.lookup path

(* Generate weak ETag from file stats: W/"mtime-size" *)
let generate_etag ~mtime ~size =
  Printf.sprintf "W/\"%x-%Lx\"" (int_of_float (mtime *. 1000.0)) size

(* Get current time *)
let now () = Unix.gettimeofday ()

(* Write common response headers *)
let write_common_headers buf ~off ~keep_alive =
  let off = Date.write_date_header buf ~off (now ()) in
  let off = Res.write_header_name buf ~off Header_name.Server "httpzo/0.1" in
  Res.write_connection buf ~off ~keep_alive

(* Write response headers for a full file response *)
let write_file_headers conn ~off status content_type file_size etag mtime version =
  let buf = conn.write_buf in
  let off = Res.write_status_line buf ~off status version in
  let off = Res.write_header_name buf ~off Header_name.Content_type content_type in
  let off = Res.write_content_length buf ~off (Int64.to_int file_size) in
  let off = Range.write_accept_ranges buf ~off in
  let off = Res.write_header buf ~off "ETag" etag in
  let off = Date.write_last_modified buf ~off mtime in
  let off = write_common_headers buf ~off ~keep_alive:conn.keep_alive in
  Res.write_crlf buf ~off

(* Write response headers for a partial content (206) response *)
let write_partial_headers conn ~off content_type ~start ~end_ ~total etag mtime version =
  let buf = conn.write_buf in
  let off = Res.write_status_line buf ~off Res.Partial_content version in
  let off = Res.write_header_name buf ~off Header_name.Content_type content_type in
  let content_length = Int64.(to_int (sub (add (sub end_ start) 1L) 0L)) in
  let off = Res.write_content_length buf ~off content_length in
  let off = Range.write_content_range buf ~off ~start ~end_ ~total in
  let off = Res.write_header buf ~off "ETag" etag in
  let off = Date.write_last_modified buf ~off mtime in
  let off = write_common_headers buf ~off ~keep_alive:conn.keep_alive in
  Res.write_crlf buf ~off

(* Write 304 Not Modified response *)
let write_not_modified_headers conn ~off etag mtime version =
  let buf = conn.write_buf in
  let off = Res.write_status_line buf ~off Res.Not_modified version in
  let off = Res.write_header buf ~off "ETag" etag in
  let off = Date.write_last_modified buf ~off mtime in
  let off = write_common_headers buf ~off ~keep_alive:conn.keep_alive in
  Res.write_crlf buf ~off

(* Write 416 Range Not Satisfiable response *)
let write_range_not_satisfiable conn ~off total version =
  let buf = conn.write_buf in
  let off = Res.write_status_line buf ~off Res.Range_not_satisfiable version in
  let off = Range.write_content_range_unsatisfiable buf ~off ~total in
  let off = write_common_headers buf ~off ~keep_alive:conn.keep_alive in
  Res.write_crlf buf ~off

(* Write buffer to flow *)
let write_buf conn ~len =
  let cs = Cstruct.of_bigarray conn.write_buf ~off:0 ~len in
  Eio.Flow.write conn.flow [cs]

(* Send error response *)
let send_error conn status message version =
  let buf = conn.write_buf in
  let off = Res.write_status_line buf ~off:0 status version in
  let off = Res.write_header_name buf ~off Header_name.Content_type "text/plain" in
  let off = Res.write_content_length buf ~off (String.length message) in
  let off = write_common_headers buf ~off ~keep_alive:conn.keep_alive in
  let off = Res.write_crlf buf ~off in
  write_buf conn ~len:off;
  Eio.Flow.write conn.flow [Cstruct.of_string message]

(* Normalize path - remove .. and resolve to absolute within root *)
let normalize_path ~root request_path =
  let decoded = request_path in
  let parts = String.split_on_char '/' decoded in
  let rec resolve acc = function
    | [] -> List.rev acc
    | "" :: rest | "." :: rest -> resolve acc rest
    | ".." :: rest ->
      (match acc with
       | [] -> resolve [] rest
       | _ :: acc' -> resolve acc' rest)
    | part :: rest -> resolve (part :: acc) rest
  in
  let normalized = resolve [] parts in
  let relative = String.concat "/" normalized in
  Eio.Path.(root / relative)

(* File metadata for caching decisions *)
type file_meta =
  { size : int64
  ; mtime : float
  ; etag : string
  ; content_type : string
  }

(* Extracted/parsed header values for conditional requests and ranges *)
type request_headers =
  { if_none_match : string option
  ; range_count : int
  }

(* Get file metadata *)
let get_file_meta file_path =
  let (_fs, path_str) = file_path in
  let stats = Unix.stat path_str in
  let size = stats.Unix.st_size |> Int64.of_int in
  let mtime = stats.Unix.st_mtime in
  let etag = generate_etag ~mtime ~size in
  let content_type = mime_type_of_path path_str in
  { size; mtime; etag; content_type }

(* Check If-None-Match header for conditional GET *)
let check_if_none_match etag if_none_match_str =
  match if_none_match_str with
  | None -> false
  | Some value ->
    if String.trim value = "*" then true
    else
      let normalize_etag s =
        let s = String.trim s in
        if String.length s >= 2 && String.sub s 0 2 = "W/" then
          String.sub s 2 (String.length s - 2)
        else s
      in
      let our_value = normalize_etag etag in
      let tags = String.split_on_char ',' value in
      List.exists (fun tag ->
        let their_value = normalize_etag tag in
        our_value = their_value
      ) tags

(* Send file with support for range requests and conditional GET *)
let send_file_with_meta conn ~file_path ~meta ~req_headers ~version =
  let { size; mtime; etag; content_type } = meta in
  let (_fs, path_str) = file_path in
  if check_if_none_match etag req_headers.if_none_match then (
    let off = write_not_modified_headers conn ~off:0 etag mtime version in
    write_buf conn ~len:off
  )
  else if req_headers.range_count = 0 then (
    (* Full content response *)
    let off = write_file_headers conn ~off:0 Res.Success content_type size etag mtime version in
    write_buf conn ~len:off;
    (* Stream file contents *)
    let fd = Unix.openfile path_str [Unix.O_RDONLY] 0 in
    Fun.protect ~finally:(fun () -> Unix.close fd) (fun () ->
      let buf = Bytes.create 65536 in
      let rec copy () =
        let n = Unix.read fd buf 0 65536 in
        if n > 0 then (
          Eio.Flow.write conn.flow [Cstruct.of_bytes buf ~off:0 ~len:n];
          copy ()
        )
      in
      copy ()
    )
  )
  else (
    (* Range request - evaluate ranges against file size *)
    let (result, _resolved_count) =
      Range.evaluate conn.ranges ~count:req_headers.range_count
        ~resource_length:size conn.resolved
    in
    match result with
    | Range.Full_content ->
      let off = write_file_headers conn ~off:0 Res.Success content_type size etag mtime version in
      write_buf conn ~len:off;
      let fd = Unix.openfile path_str [Unix.O_RDONLY] 0 in
      Fun.protect ~finally:(fun () -> Unix.close fd) (fun () ->
        let buf = Bytes.create 65536 in
        let rec copy () =
          let n = Unix.read fd buf 0 65536 in
          if n > 0 then (
            Eio.Flow.write conn.flow [Cstruct.of_bytes buf ~off:0 ~len:n];
            copy ()
          )
        in
        copy ()
      )
    | Range.Not_satisfiable ->
      conn.keep_alive <- false;
      let off = write_range_not_satisfiable conn ~off:0 size version in
      write_buf conn ~len:off
    | Range.Single_range | Range.Multiple_ranges ->
      let r = conn.resolved.(0) in
      let start = r.start in
      let end_ = r.end_ in
      let range_len = Int64.(sub (add (sub end_ start) 1L) 0L) in
      let len = Int64.to_int range_len in
      let off = write_partial_headers conn ~off:0 content_type ~start ~end_ ~total:size etag mtime version in
      write_buf conn ~len:off;
      let fd = Unix.openfile path_str [Unix.O_RDONLY] 0 in
      Fun.protect ~finally:(fun () -> Unix.close fd) (fun () ->
        let _ = Unix.lseek fd (Int64.to_int start) Unix.SEEK_SET in
        let buf = Bytes.create (min len 65536) in
        let remaining = ref len in
        while !remaining > 0 do
          let to_read = min !remaining 65536 in
          let n = Unix.read fd buf 0 to_read in
          if n > 0 then (
            Eio.Flow.write conn.flow [Cstruct.of_bytes buf ~off:0 ~len:n];
            remaining := !remaining - n
          ) else
            remaining := 0
        done
      )
  )

(* Try to serve index.html from a directory *)
let serve_directory conn ~file_path ~req_headers ~version =
  let index_path = Eio.Path.(file_path / "index.html") in
  let (_fs, index_str) = index_path in
  if Sys.file_exists index_str then (
    let meta = get_file_meta index_path in
    send_file_with_meta conn ~file_path:index_path ~meta ~req_headers ~version
  ) else
    send_error conn Res.Not_found "Not Found" version

(* Try to serve a regular file, checking it's within root *)
let serve_regular_file conn ~root ~file_path ~req_headers ~version =
  let (_root_fs, root_str) = root in
  let (_fs, path_str) = file_path in
  try
    let file_abs = Unix.realpath path_str in
    let root_abs = Unix.realpath root_str in
    if String.length file_abs >= String.length root_abs &&
       String.sub file_abs 0 (String.length root_abs) = root_abs
    then (
      let meta = get_file_meta file_path in
      send_file_with_meta conn ~file_path ~meta ~req_headers ~version
    ) else
      send_error conn Res.Forbidden "Forbidden" version
  with _ ->
    send_error conn Res.Not_found "Not Found" version

(* Serve a file *)
let serve_file conn ~root target_str req_headers version =
  let path = match String.index_opt target_str '?' with
    | Some idx -> String.sub target_str 0 idx
    | None -> target_str
  in
  let file_path = normalize_path ~root path in
  let (_fs, path_str) = file_path in
  if Sys.file_exists path_str then (
    if Sys.is_directory path_str then
      serve_directory conn ~file_path ~req_headers ~version
    else
      serve_regular_file conn ~root ~file_path ~req_headers ~version
  ) else
    send_error conn Res.Not_found "Not Found" version

(* Read more data into buffer *)
let read_more conn =
  if conn.read_len >= Buf_read.buffer_size then
    `Buffer_full
  else (
    let available = Buf_read.buffer_size - conn.read_len in
    let cs = Cstruct.of_bigarray conn.read_buf ~off:conn.read_len ~len:available in
    match Eio.Flow.single_read conn.flow cs with
    | n ->
      conn.read_len <- conn.read_len + n;
      `Ok n
    | exception End_of_file -> `Eof
  )

(* Shift buffer contents to remove processed data *)
let shift_buffer conn consumed =
  if consumed > 0 && consumed < conn.read_len then (
    for i = 0 to conn.read_len - consumed - 1 do
      Bigarray.Array1.set conn.read_buf i
        (Bigarray.Array1.get conn.read_buf (consumed + i))
    done;
    conn.read_len <- conn.read_len - consumed
  ) else if consumed >= conn.read_len then
    conn.read_len <- 0

(* Handle one request on connection *)
let handle_request conn ~root ~limits =
  let buf = conn.read_buf in
  let len = conn.read_len in
  let (status, req, headers) = Request_parse.parse buf ~len ~limits in
  let body_off = req.Req.body_off in
  let version = req.Req.version in
  let target = req.Req.target in
  match status with
  | Buf_read.Complete ->
    let target_str = Span.to_string buf target in
    let if_none_match =
      match Header.find headers Header_name.If_none_match with
      | None -> None
      | Some hdr -> Some (Span.to_string buf hdr.Header.value)
    in
    let range_count =
      match Header.find headers Header_name.Range with
      | None -> 0
      | Some hdr ->
        let (status, count) = Range.parse buf hdr.Header.value conn.ranges in
        match status with
        | Range.Invalid -> 0
        | Range.Valid -> count
    in
    let req_headers = { if_none_match; range_count } in
    let body_span = Req.body_span ~len req in
    let body_span_len = Span.len body_span in
    let body_span_off = Span.off body_span in
    if body_span_len = -1 then
      `Need_more
    else (
      conn.keep_alive <- req.Req.keep_alive;
      serve_file conn ~root target_str req_headers version;
      let consumed =
        if body_span_len > 0 then body_span_off + body_span_len else body_off
      in
      shift_buffer conn consumed;
      if conn.keep_alive then `Continue else `Close
    )
  | Buf_read.Partial -> `Need_more
  | Buf_read.Headers_too_large
  | Buf_read.Content_length_overflow ->
    conn.keep_alive <- false;
    send_error conn Res.Payload_too_large "Payload Too Large" Version.Http_1_1;
    `Close
  | Buf_read.Bare_cr_detected
  | Buf_read.Ambiguous_framing ->
    conn.keep_alive <- false;
    send_error conn Res.Bad_request "Bad Request" Version.Http_1_1;
    `Close
  | Buf_read.Missing_host_header ->
    conn.keep_alive <- false;
    send_error conn Res.Bad_request "Missing Host Header" Version.Http_1_1;
    `Close
  | _ ->
    conn.keep_alive <- false;
    send_error conn Res.Bad_request "Bad Request" Version.Http_1_1;
    `Close

(* Send payload too large error and close connection *)
let send_payload_too_large conn =
  conn.keep_alive <- false;
  send_error conn Res.Payload_too_large "Payload Too Large" Version.Http_1_1

(* Handle connection loop *)
let handle_connection conn ~root ~limits =
  let handle_read_result ~continue = function
    | `Eof -> ()
    | `Buffer_full -> send_payload_too_large conn
    | `Ok _ -> continue ()
  in
  let rec loop () =
    if conn.read_len = 0 then
      handle_read_result ~continue:loop (read_more conn)
    else
      match handle_request conn ~root ~limits with
      | `Continue -> loop ()
      | `Close -> ()
      | `Need_more -> handle_read_result ~continue:loop (read_more conn)
  in
  loop ()

(* Handle a single client connection *)
let handle_client ~root ~limits flow _addr =
  let conn = create_conn flow in
  try
    handle_connection conn ~root ~limits
  with exn ->
    traceln "Error: %s" (Printexc.to_string exn)

(* Run server *)
let run ~net ~sw config =
  traceln "httpzo serving on http://localhost:%d/" config.port;
  traceln "  Supports: Range requests, ETag, If-None-Match";
  let addr = `Tcp (Eio.Net.Ipaddr.V4.loopback, config.port) in
  let sock = Eio.Net.listen net ~sw ~backlog:128 ~reuse_addr:true addr in
  let rec accept_loop () =
    Eio.Net.accept_fork sock ~sw ~on_error:(fun exn ->
      traceln "Connection error: %s" (Printexc.to_string exn)
    ) (fun flow addr -> handle_client ~root:config.root ~limits:config.limits flow addr);
    accept_loop ()
  in
  accept_loop ()