(** Chunk file management for partial downloads *)

open Eio

module Range = Range
module Flags = Flags

type t = {
  key : string;
  hash : string;
  range : Range.t;
  path : Fs.dir_ty Path.t;
  flags : Flags.t;
  ttl : float option;
}

let key c = c.key
let hash c = c.hash
let range c = c.range
let path c = c.path
let flags c = c.flags
let ttl c = c.ttl


(** Parse a chunk from a filename *)
let parse ~path ~filename =
  try
    (* Expected format: {hash_rest}:2:{ttl},{flags}:{start}-{end} *)
    match String.split_on_char ':' filename with
    | hash_rest :: "2" :: rest when List.length rest >= 1 ->
        (* Extract components *)
        let ttl_flags_range = String.concat ":" rest in

        (* Find the last colon that separates range *)
        (match String.rindex_opt ttl_flags_range ':' with
        | None -> None  (* No range part *)
        | Some idx ->
            let ttl_flags = String.sub ttl_flags_range 0 idx in
            let range_str = String.sub ttl_flags_range (idx + 1)
                          (String.length ttl_flags_range - idx - 1) in

            (* Parse range *)
            match Range.of_string range_str with
            | None -> None
            | Some range ->
                (* Parse TTL and flags *)
                let ttl, flags_str =
                  match String.split_on_char ',' ttl_flags with
                  | [] -> None, ""
                  | [ttl_or_flags] ->
                      if String.contains ttl_or_flags ',' then
                        None, ttl_or_flags
                      else
                        (try Some (float_of_string ttl_or_flags) with _ -> None), ""
                  | ttl :: rest ->
                      (try Some (float_of_string ttl) with _ -> None),
                      String.concat "," rest
                in

                let flags = Flags.of_string flags_str in

                (* Only consider it a chunk if it has the Chunk flag *)
                if Flags.is_chunk flags then
                  (* We don't have the full hash or key from just the filename *)
                  (* The caller needs to provide these through context *)
                  Some {
                    key = "";  (* Will be filled by caller *)
                    hash = hash_rest;  (* Partial hash *)
                    range;
                    path;
                    flags;
                    ttl;
                  }
                else
                  None)
    | _ -> None
  with _ -> None

(** Generate a chunk filename *)
let make_filename ~hash ~range ?ttl ?(flags=Flags.empty) () =
  (* Ensure Chunk flag is set *)
  let flags = Flags.add `Chunk flags in

  let hash_rest = String.sub hash 4 (String.length hash - 4) in
  let ttl_part = match ttl with
    | None -> ""
    | Some t -> Printf.sprintf "%d" (int_of_float t)
  in
  let flags_str = Flags.to_string flags in
  let range_str = Range.to_string range in

  Printf.sprintf "%s:2:%s,%s:%s" hash_rest ttl_part flags_str range_str

(** Hash a key using SHA256 *)
let hash_key key =
  Digestif.SHA256.(to_hex (digest_string key))

(** Get directory components from hash *)
let key_to_dirs key =
  let hash = hash_key key in
  if String.length hash < 4 then
    invalid_arg "hash too short"
  else
    [ String.sub hash 0 2;
      String.sub hash 2 2 ]

(** Find all chunks for a key *)
let find_chunks ~base_dir ~key =
  let dirs = key_to_dirs key in
  let dir_path = List.fold_left (fun acc d -> Path.(acc / d)) base_dir dirs in

  try
    let entries = Path.read_dir dir_path in
    let hash = hash_key key in
    let hash_rest = String.sub hash 4 (String.length hash - 4) in

    List.filter_map (fun entry ->
      if String.starts_with ~prefix:(hash_rest ^ ":2") entry then
        match parse ~path:Path.(dir_path / entry) ~filename:entry with
        | Some chunk ->
            (* Fill in the complete key and hash *)
            Some { chunk with key; hash }
        | None -> None
      else None
    ) entries
  with _ -> []

(** Sort chunks by range start position *)
let sort_by_range chunks =
  List.sort (fun c1 c2 -> Range.compare c1.range c2.range) chunks

(** Check if chunks form a complete sequence *)
let is_complete_sequence chunks ~total_size =
  let sorted = sort_by_range chunks in

  let rec check_continuous pos = function
    | [] -> pos >= total_size
    | chunk :: rest ->
        if Range.start chunk.range > pos then
          false  (* Gap found *)
        else
          let next_pos = Int64.succ (Range.end_ chunk.range) in
          check_continuous (max pos next_pos) rest
  in

  check_continuous 0L sorted

(** Get missing ranges from a list of chunks *)
let missing_ranges chunks ~total_size =
  let sorted = sort_by_range chunks in

  let rec find_gaps acc pos = function
    | [] ->
        if pos < total_size then
          (* Gap at the end *)
          List.rev (Range.create ~start:pos ~end_:(Int64.pred total_size) :: acc)
        else
          List.rev acc
    | chunk :: rest ->
        let chunk_start = Range.start chunk.range in
        if chunk_start > pos then
          (* Found a gap *)
          let gap = Range.create ~start:pos ~end_:(Int64.pred chunk_start) in
          find_gaps (gap :: acc) (Int64.succ (Range.end_ chunk.range)) rest
        else
          (* No gap, move to next position *)
          find_gaps acc (max pos (Int64.succ (Range.end_ chunk.range))) rest
  in

  find_gaps [] 0L sorted

let pp fmt c =
  Format.fprintf fmt "Chunk{key=%s, range=%a, flags=%a}"
    c.key Range.pp c.range Flags.pp c.flags