(* Copyright (c) 2024-2026 Thomas Gazagnaire <thomas@gazagnaire.org>

   Permission to use, copy, modify, and distribute this software for any
   purpose with or without fee is hereby granted, provided that the above
   copyright notice and this permission notice appear in all copies.

   THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
   WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
   MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
   ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
   WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
   ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
   OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. *)

type entry = { path : Fpath.t; head : Hash.t; branch : string option }
type t = { fs : Eio.Fs.dir_ty Eio.Path.t; git_dir : string }

let v ~fs ~git_dir = { fs; git_dir }

let err_worktree_not_found path =
  Error (`Msg (Fmt.str "worktree '%a' not found" Fpath.pp path))

(* Get the main worktree path from .git directory *)
let main_worktree_path t =
  (* .git dir is like /path/to/repo/.git, so parent is the worktree *)
  if
    String.length t.git_dir > 5
    && String.sub t.git_dir (String.length t.git_dir - 5) 5 = "/.git"
  then String.sub t.git_dir 0 (String.length t.git_dir - 5)
  else Filename.dirname t.git_dir

(* Get worktree name from path *)
let name path = Fpath.basename path

(* Parse HEAD file to get branch name *)
let parse_head_file content =
  let content = String.trim content in
  if String.length content > 16 && String.sub content 0 16 = "ref: refs/heads/"
  then Some (String.sub content 16 (String.length content - 16))
  else None

(* Read linked worktree entry from .git/worktrees/<name> *)
let read_worktree_entry t name =
  let wt_git_dir = Filename.concat t.git_dir ("worktrees/" ^ name) in
  let gitdir_path = Eio.Path.(t.fs / wt_git_dir / "gitdir") in
  let head_path = Eio.Path.(t.fs / wt_git_dir / "HEAD") in
  try
    let gitdir_content = String.trim (Eio.Path.load gitdir_path) in
    (* gitdir contains path to the worktree's .git file, the worktree is the parent *)
    let wt_path =
      if Filename.check_suffix gitdir_content "/.git" then
        String.sub gitdir_content 0 (String.length gitdir_content - 5)
      else Filename.dirname gitdir_content
    in
    let head_content = Eio.Path.load head_path in
    let branch = parse_head_file head_content in
    let head_hash =
      match branch with
      | Some b -> (
          (* Read from the shared refs *)
          let ref_path = Eio.Path.(t.fs / t.git_dir / "refs" / "heads" / b) in
          try
            let hash_str = String.trim (Eio.Path.load ref_path) in
            Hash.of_hex hash_str
          with Eio.Io _ | Invalid_argument _ ->
            (* Fallback: parse HEAD as direct hash *)
            Hash.of_hex (String.trim head_content))
      | None -> Hash.of_hex (String.trim head_content)
    in
    match Fpath.of_string wt_path with
    | Ok path -> Some { path; head = head_hash; branch }
    | Error _ -> None
  with Eio.Io _ | Invalid_argument _ -> None

let list t ~head ~current_branch =
  (* First, add the main worktree *)
  let main_path = main_worktree_path t in
  let main_entry =
    match Fpath.of_string main_path with
    | Ok path ->
        (* For the main worktree, use a zero hash if HEAD doesn't exist yet *)
        let head_hash =
          match head with
          | Some h -> h
          | None -> Hash.of_hex (String.make 40 '0')
        in
        Some { path; head = head_hash; branch = current_branch }
    | Error _ -> None
  in
  (* Then list linked worktrees from .git/worktrees/ *)
  let worktrees_dir = Filename.concat t.git_dir "worktrees" in
  let worktrees_path = Eio.Path.(t.fs / worktrees_dir) in
  let linked_entries =
    try
      let entries = Eio.Path.read_dir worktrees_path in
      List.filter_map (read_worktree_entry t) entries
    with Eio.Io _ -> []
  in
  match main_entry with Some e -> e :: linked_entries | None -> linked_entries

let exists t ~path =
  (* We need to get head and current_branch, but for exists check we can pass None *)
  let worktrees = list t ~head:None ~current_branch:None in
  List.exists (fun e -> Fpath.equal e.path path) worktrees

let write_ref t name hash =
  let path = Filename.concat t.git_dir name in
  let full_path = Eio.Path.(t.fs / path) in
  let dir = Filename.dirname path in
  let dir_path = Eio.Path.(t.fs / dir) in
  (try Eio.Path.mkdir ~perm:0o755 dir_path with Eio.Io _ -> ());
  let content = Hash.to_hex hash ^ "\n" in
  Eio.Path.save ~create:(`Or_truncate 0o644) full_path content

let add t ~head ~path ~branch =
  let name = name path in
  let wt_git_dir = Filename.concat t.git_dir ("worktrees/" ^ name) in
  let wt_git_dir_path = Eio.Path.(t.fs / wt_git_dir) in
  let wt_path_str = Fpath.to_string path in
  let wt_path = Eio.Path.(t.fs / wt_path_str) in
  (* Create the worktree directory *)
  (try Eio.Path.mkdirs ~perm:0o755 wt_path with Eio.Io _ -> ());
  (* Create .git/worktrees/<name> directory *)
  (try Eio.Path.mkdirs ~perm:0o755 wt_git_dir_path with Eio.Io _ -> ());
  (* Write gitdir file (path to the worktree's .git file) *)
  let gitdir_content = wt_path_str ^ "/.git\n" in
  Eio.Path.save ~create:(`Or_truncate 0o644)
    Eio.Path.(wt_git_dir_path / "gitdir")
    gitdir_content;
  (* Write HEAD file (pointing to the new branch) *)
  let head_content = "ref: refs/heads/" ^ branch ^ "\n" in
  Eio.Path.save ~create:(`Or_truncate 0o644)
    Eio.Path.(wt_git_dir_path / "HEAD")
    head_content;
  (* Write commondir file (relative path to main .git) *)
  Eio.Path.save ~create:(`Or_truncate 0o644)
    Eio.Path.(wt_git_dir_path / "commondir")
    "..\n";
  (* Create the branch in the main repo *)
  write_ref t ("refs/heads/" ^ branch) head;
  (* Write .git file in the worktree *)
  let git_file_content = "gitdir: " ^ wt_git_dir ^ "\n" in
  Eio.Path.save ~create:(`Or_truncate 0o644)
    Eio.Path.(wt_path / ".git")
    git_file_content;
  Ok ()

let remove t ~path ~force =
  if not (exists t ~path) then err_worktree_not_found path
  else
    let name = name path in
    let wt_git_dir = Filename.concat t.git_dir ("worktrees/" ^ name) in
    let wt_git_dir_path = Eio.Path.(t.fs / wt_git_dir) in
    let wt_path_str = Fpath.to_string path in
    let wt_path = Eio.Path.(t.fs / wt_path_str) in
    (* Check if worktree is the main worktree *)
    if String.equal (main_worktree_path t) wt_path_str then
      Error (`Msg "cannot remove main worktree")
    else begin
      (* TODO: Check for uncommitted changes if not force *)
      ignore force;
      (* Remove the .git/worktrees/<name> directory *)
      let rec remove_dir path =
        try
          let entries = Eio.Path.read_dir path in
          List.iter
            (fun entry ->
              let entry_path = Eio.Path.(path / entry) in
              if Eio.Path.is_directory entry_path then remove_dir entry_path
              else Eio.Path.unlink entry_path)
            entries;
          Eio.Path.rmdir path
        with Eio.Io _ -> ()
      in
      remove_dir wt_git_dir_path;
      (* Remove the worktree directory *)
      remove_dir wt_path;
      Ok ()
    end