include Disk_sig.Types

module Hashtbl_base = Hashtbl
module Hashtbl = struct
  include Hashtbl_base
  let find_opt t x = try Some (find t x) with
  | Not_found ->
    None

  let empty t = (length t) = 0
end;;


type file =
  | Actual_file_with_contents of string
  | Directory of (string, file) Hashtbl.t

module Helpers = struct
  exception Relative_parent_not_supported
  exception Is_not_actual_file
  exception Cannot_overwrite_existing_directory_with_actual_file

end

(** Directory for "/" *)
let root = Hashtbl.create 10;;

(** We avoid using Unix.getcwd () in TestDisk because:
  *   1) Getting global state from this clean test environment is gross
  *   2) Because of 1, CWD should actually be tracked inside this virtual FS.
  *   3) Javascript doesn't support Unix.getcwd anyway, so we can't transpile
  * As such, we make a fake cwd *)
let cwd = ref "/fake/initial_cwd"

let getcwd () = !cwd
let chdir s = cwd := s

let rec mkdir_p path root = match Filename.dirname path with
  | "." ->
    mkdir_p (getcwd ()) root
  | "/" ->
    root
  | parent ->
    let parent = mkdir_p parent root in
    if (Filename.basename path) = "." then
      parent
    else if (Filename.basename path) = ".." then
      raise Helpers.Relative_parent_not_supported
    else begin
      match Hashtbl.find_opt parent (Filename.basename path) with
      | None ->
        let dir = Hashtbl.create 10 in
        let () = Hashtbl.add parent (Filename.basename path) (Directory dir) in
        dir
      | Some (Directory table) ->
        table
      | Some (Actual_file_with_contents _) ->
        raise (NotADirectory path)
    end

and get_dir path root = match Filename.dirname path with
  | "." ->
    get_dir (getcwd ()) root
  | "/" ->
    root
  | parent ->
    let parent = get_dir parent root in
    if (Filename.basename path) = "." then
      parent
    else if (Filename.basename path) = ".." then
      raise Helpers.Relative_parent_not_supported
    else begin
      match Hashtbl.find_opt parent (Filename.basename path) with
      | None ->
        raise (No_such_file_or_directory path)
      | Some (Directory table) ->
        table
      | Some (Actual_file_with_contents _) ->
        raise (NotADirectory path)
    end

(** Returns file at path (may be an actual file or a directory). *)
and get_file path root =
  let parent = get_dir (Filename.dirname path) root in
  let basename = Filename.basename path in
  if basename = "." then
    (Directory parent)
  else begin
    try Hashtbl.find parent basename with
    | Not_found ->
      raise (No_such_file_or_directory path)
  end

(** Initialize creation of CWD. *)
let () = ignore (mkdir_p "." root)

let get x =
  match get_file x root with
  | Actual_file_with_contents contents ->
    contents
  | Directory _ ->
    raise Helpers.Is_not_actual_file


(** Set the contents "y" for file "x". Has an option to create all parent
 * directories automatically. *)
let set ?(create_parent_dirs=true) x y =
  let parent = if create_parent_dirs then
    mkdir_p (Filename.dirname x) root
  else
    get_dir (Filename.dirname x) root
  in
  match Hashtbl.find_opt parent (Filename.basename x) with
  | None ->
    let new_file = Actual_file_with_contents y in
    Hashtbl.add parent (Filename.basename x) new_file
  | Some (Actual_file_with_contents _) ->
    let new_file = Actual_file_with_contents y in
    Hashtbl.replace parent (Filename.basename x) new_file
  | Some (Directory _) ->
    raise Helpers.Cannot_overwrite_existing_directory_with_actual_file

let is_directory x =
  try begin
    match get_file x root with
    | Directory _ ->
      true
    | Actual_file_with_contents _ ->
      false
  end
  with
  | No_such_file_or_directory _ ->
    false

let cat = get

let file_exists x =
  try
    match get_file x root with
    | Actual_file_with_contents _
    | Directory _ ->
      true
  with
  | No_such_file_or_directory _ ->
    false

let write_file ~file ~contents = set ~create_parent_dirs:false file contents

let mkdir path _perm =
  let parent = get_dir (Filename.dirname path) root in
  ignore (mkdir_p (Filename.basename path) parent)

let mkdir_p path = ignore (mkdir_p path root)

let rm_dir_tree path =
  if path = "/" then
    Hashtbl.clear root
  else try
    let dir = get_dir (Filename.dirname path) root in
    Hashtbl.remove dir (Filename.basename path)
  with
  | No_such_file_or_directory _ ->
    (** File already doesn't exist; ignore. *)
    ()

let readdir x = match get_file x root with
  | Actual_file_with_contents _ ->
    raise (NotADirectory x)
  | Directory directory ->
    let names = Hashtbl.fold (fun k _v acc -> k :: acc) directory [] in
    Array.of_list names

let rename old target =
  if not (file_exists old) then
    raise (No_such_file_or_directory old)
  else if not (file_exists (Filename.dirname target)) then
    raise (No_such_file_or_directory (Filename.dirname target))
  else
    let old_parent = get_dir (Filename.dirname old) root in
    let old_file = get_file old root in
    (** What if the last character in target is a "/"? What to do? *)
    let target_parent = get_dir (Filename.dirname target) root in
    match old_file, Hashtbl.find_opt target_parent (Filename.basename target) with
    | Directory _, Some (Directory target_files) when not (Hashtbl.empty target_files) ->
      raise (Rename_target_dir_not_empty target)
    | Directory _, Some (Directory _)
    | _, None ->
      (** Rename one directory to the other. *)
      Hashtbl.replace target_parent (Filename.basename target) old_file;
      Hashtbl.remove old_parent (Filename.basename old);
    | _, _ ->
      failwith "Not sure what to do here"
