(*-----------------------------------------------------------------
   Copyright (c) 2025 The merry programmers. All rights reserved.
   SPDX-License-Identifier: ISC
  -----------------------------------------------------------------*)
open Import
open Exit.Syntax

module Options = struct
  type t = { noclobber : bool }

  let default = { noclobber = false }

  let with_options ?noclobber t =
    { noclobber = Option.value ~default:t.noclobber noclobber }
end

(** An evaluator over the AST *)
module Make (S : Types.State) (E : Types.Exec) = struct
  (* What follows uses the POSIX definition of what a shell does ($ 2.1).

   It starts from point (4), completing a series of expansions on the AST,
   then redirection is setup, and finally functions/built-ins/commands are
   executed. *)

  class default_map =
    object (_)
      inherit Ast.map
      method string (s : string) = s
      method int (i : int) = i
      method char c = c
      method option f v = Option.map f v
      method nlist__t f t = Nlist.map f t
      method nslist__t f t = Nslist.map f t
      method list f t = List.map f t
    end

  type ctx = {
    interactive : bool;
    state : S.t;
    local_state : (string * string) list;
    executor : E.t;
    fs : Eio.Fs.dir_ty Eio.Path.t;
    options : Options.t;
    stdout : Eio_unix.sink_ty Eio.Flow.sink option;
  }

  let clear_local_state ctx = { ctx with local_state = [] }

  class default_ctx_fold =
    object (_)
      inherit [ctx] Ast.fold
      method int _ ctx = ctx
      method string _ ctx = ctx
      method char _ ctx = ctx
      method option f v ctx = Option.fold ~none:ctx ~some:(fun i -> f i ctx) v
      method nlist__t f v ctx = Nlist.fold_left (fun acc i -> f i acc) ctx v

      method nslist__t f g v ctx =
        Nslist.fold_left (fun acc a b -> f a acc |> g b) ctx v

      method list f v ctx = List.fold_left (fun acc i -> f i acc) ctx v
    end

  let map_word_components f ast =
    let o =
      object (_)
        inherit default_map
        method! word_component cst = f cst
      end
    in
    o#complete_command ast

  let map_words ?(skip_for_clauses = true) f =
    let o =
      object (_)
        inherit default_map as super
        method! word cst = f cst

        method! for_clause cst =
          if skip_for_clauses then cst else super#for_clause cst
      end
    in
    o

  let rec tilde_expansion ctx = function
    | [] -> []
    | Ast.WordTildePrefix _ :: rest ->
        Ast.WordName (S.expand ctx.state `Tilde) :: tilde_expansion ctx rest
    | v :: rest -> v :: tilde_expansion ctx rest

  let parameter_expansion' ctx =
    let rec expand = function
      | [] -> []
      | Ast.WordVariable v :: rest -> (
          match v with
          | Ast.VariableAtom (s, NoAttribute) -> (
              match S.lookup ctx.state ~param:s with
              | None -> Ast.WordName "" :: expand rest
              | Some cst -> cst @ expand rest)
          | _ -> Fmt.failwith "No support for variable attributes yet!")
      | Ast.WordDoubleQuoted cst :: rest ->
          Ast.WordDoubleQuoted (expand cst) :: expand rest
      | Ast.WordSingleQuoted cst :: rest ->
          Ast.WordSingleQuoted (expand cst) :: expand rest
      | v :: rest -> v :: expand rest
    in
    (ctx, expand)

  let stdout_for_pipeline ~sw ctx = function
    | [] -> (None, ctx.stdout)
    | _ ->
        let r, w = Eio_unix.pipe sw in
        (Some r, Some (w :> Eio_unix.sink_ty Eio.Flow.sink))

  let fd_of_int ?(close_unix = true) ~sw n =
    Eio_unix.Fd.of_unix ~close_unix ~sw (Obj.magic n : Unix.file_descr)

  let handle_one_redirection ~sw ctx = function
    | Ast.IoRedirect_IoFile (n, (op, file)) -> (
        match op with
        | Io_op_less ->
            (* Simple redirection for input *)
            let r =
              Eio.Path.open_in ~sw (ctx.fs / Ast.word_components_to_string file)
            in
            let fd = Eio_unix.Resource.fd_opt r |> Option.get in
            Some (Types.Redirect (n, fd, `Blocking))
        | Io_op_lessand -> (
            match file with
            | [ WordLiteral "-" ] ->
                if n = 0 then Some (Types.Close Eio_unix.Fd.stdin)
                else
                  let fd = fd_of_int ~sw n in
                  Some (Types.Close fd)
            | [ WordLiteral m ] when Option.is_some (int_of_string_opt m) ->
                let m = int_of_string m in
                Some
                  (Types.Redirect
                     (n, fd_of_int ~close_unix:false ~sw m, `Blocking))
            | _ -> None)
        | (Io_op_great | Io_op_dgreat) as v ->
            (* Simple file creation *)
            let append = v = Io_op_dgreat in
            let w =
              Eio.Path.open_out ~sw ~append ~create:(`If_missing 0o644)
                (ctx.fs / Ast.word_components_to_string file)
            in
            let fd = Eio_unix.Resource.fd_opt w |> Option.get in
            Some (Types.Redirect (n, fd, `Blocking))
        | Io_op_greatand -> (
            match file with
            | [ WordLiteral "-" ] ->
                if n = 0 then Some (Types.Close Eio_unix.Fd.stdin)
                else
                  let fd = fd_of_int ~sw n in
                  Some (Types.Close fd)
            | [ WordLiteral m ] when Option.is_some (int_of_string_opt m) ->
                let m = int_of_string m in
                Some
                  (Types.Redirect
                     (n, fd_of_int ~close_unix:false ~sw m, `Blocking))
            | _ -> None)
        | Io_op_clobber -> Fmt.failwith ">| not supported yet."
        | Io_op_lessgreat -> Fmt.failwith "<> not support yet.")
    | Ast.IoRedirect_IoHere _ ->
        Fmt.failwith "HERE documents not yet implemented!"

  let handle_built_in (ctx : ctx) = function
    | Built_ins.Cd { path } ->
        let cwd = S.cwd ctx.state in
        let+ state =
          match path with
          | Some p ->
              let fp = Fpath.append cwd (Fpath.v p) in
              Exit.map' (Eunix.chdir p)
                ~zero:(fun () -> S.set_cwd ctx.state fp)
                ~nonzero:(fun () -> ctx.state)
          | None -> (
              match Eunix.find_env "HOME" with
              | None -> Exit.nonzero_msg ctx.state "HOME not set"
              | Some p -> Exit.zero (S.set_cwd ctx.state @@ Fpath.v p))
        in
        { ctx with state }
    | Pwd ->
        Fmt.pr "%s\n%!" (Eunix.cwd ());
        Exit.zero ctx
    | Exit n ->
        let should_exit =
          { Exit.default_should_exit with interactive = `Yes }
        in
        Exit.nonzero_msg ~should_exit ctx ~exit_code:n "exit"

  let cwd_of_ctx ctx = S.cwd ctx.state |> Fpath.to_string |> ( / ) ctx.fs

  let needs_glob_expansion : Ast.word_component -> bool = function
    | WordGlobAll | WordGlobAny -> true
    | _ -> false

  let apply_pair (a, b) f = f a b
  let ( ||> ) = apply_pair

  let get_env ?(extra = []) () =
    let env = Eunix.env () in
    List.fold_left (fun acc (k, _) -> List.remove_assoc k acc) env extra
    |> List.append extra

  let rec execute_commands initial_ctx local_switch p =
    let rec loop (exit_ctx : ctx Exit.t)
        (stdout_of_previous : Eio_unix.source_ty Eio_unix.source option) :
        Ast.command list -> ctx Exit.t =
     fun c ->
      let ctx = Exit.value exit_ctx in
      match c with
      | Ast.SimpleCommand (Prefixed (prefix, None, _suffix)) :: rest ->
          let ctx = collect_assignments ctx prefix in
          loop (Exit.zero ctx) stdout_of_previous rest
      | Ast.SimpleCommand (Prefixed (prefix, Some executable, suffix)) :: rest
        ->
          let ctx = collect_assignments ~update:false ctx prefix in
          loop (Exit.zero ctx) stdout_of_previous
            (Ast.SimpleCommand (Named (executable, suffix)) :: rest)
      | Ast.SimpleCommand (Named (executable, None)) :: rest -> (
          let ctx, executable = expand_cst ctx executable in
          match
            Built_ins.of_args
              [ handle_word_components_to_string ctx executable ]
          with
          | Some bi -> handle_built_in ctx bi
          | None -> (
              let some_read, some_write =
                stdout_for_pipeline ctx ~sw:local_switch rest
              in
              match stdout_of_previous with
              | None ->
                  let executable =
                    handle_word_components_to_string ctx executable
                  in
                  let res =
                    E.exec ctx.executor ?stdout:some_write ~cwd:(cwd_of_ctx ctx)
                      ~env:(get_env ~extra:ctx.local_state ())
                      [ executable ]
                    >|= fun () -> clear_local_state ctx
                  in
                  Option.iter Eio.Flow.close some_write;
                  loop res some_read rest
              | Some stdout ->
                  let executable =
                    handle_word_components_to_string ctx executable
                  in
                  let res =
                    E.exec ctx.executor ~stdin:stdout ?stdout:some_write
                      ~env:(get_env ~extra:ctx.local_state ())
                      ~cwd:(cwd_of_ctx ctx) [ executable ]
                    >|= fun () -> clear_local_state ctx
                  in
                  Option.iter Eio.Flow.close some_write;
                  loop res some_read rest))
      | Ast.SimpleCommand (Named (executable, Some suffix)) :: rest -> (
          let ctx, executable = expand_cst ctx executable in
          let ctx, suffix = expand_redirects (ctx, []) suffix in
          let args = args ctx suffix in
          match
            Built_ins.of_args
              (handle_word_components_to_string ctx executable :: args)
          with
          | Some bi -> handle_built_in ctx bi
          | None -> (
              let redirect =
                List.fold_left
                  (fun acc -> function
                    | Ast.Suffix_word _ -> acc
                    | Ast.Suffix_redirect rdr ->
                        handle_one_redirection ~sw:local_switch ctx rdr :: acc)
                  [] suffix
                |> List.rev |> List.filter_map Fun.id
              in
              let some_read, some_write =
                stdout_for_pipeline ~sw:local_switch ctx rest
              in
              match stdout_of_previous with
              | None ->
                  let res =
                    E.exec ~fds:redirect ctx.executor ?stdout:some_write
                      ~cwd:(cwd_of_ctx ctx)
                      ~env:(get_env ~extra:ctx.local_state ())
                      (handle_word_components_to_string ctx executable :: args)
                    >|= fun () -> clear_local_state ctx
                  in
                  Option.iter Eio.Flow.close some_write;
                  loop res some_read rest
              | Some stdout ->
                  let res =
                    E.exec ~fds:redirect ctx.executor ~stdin:stdout
                      ~cwd:(cwd_of_ctx ctx) ?stdout:some_write
                      ~env:(get_env ~extra:ctx.local_state ())
                      (handle_word_components_to_string ctx executable :: args)
                    >|= fun () -> clear_local_state ctx
                  in
                  Option.iter Eio.Flow.close some_write;
                  loop res some_read rest))
      | CompoundCommand (c, rdrs) :: _rest ->
          let _rdrs =
            List.map (handle_one_redirection ~sw:local_switch ctx) rdrs
          in
          let ctx = handle_compound_command ctx c in
          ctx
      | v :: _ ->
          Fmt.epr "TODO: %a" Yojson.Safe.pp (Ast.command_to_yojson v);
          failwith "Err"
      | [] -> exit_ctx
    in
    loop (Exit.zero initial_ctx) None p

  and expand_cst (ctx : ctx) cst =
    let cst = tilde_expansion ctx cst in
    let _, o = parameter_expansion' ctx in
    (ctx, o cst)

  and expand_redirects ((ctx, acc) : ctx * Ast.cmd_suffix_item list)
      (c : Ast.cmd_suffix_item list) =
    match c with
    | [] -> (ctx, List.rev acc)
    | Ast.Suffix_redirect (IoRedirect_IoFile (num, (op, file))) :: rest ->
        let ctx, cst = expand_cst ctx file in
        let cst = handle_subshell ctx cst in
        let v = Ast.Suffix_redirect (IoRedirect_IoFile (num, (op, cst))) in
        expand_redirects (ctx, v :: acc) rest
    | (Ast.Suffix_redirect _ as v) :: rest ->
        expand_redirects (ctx, v :: acc) rest
    | s :: rest -> expand_redirects (ctx, s :: acc) rest

  and handle_single_pipeline ~sw ctx c =
    let pipeline = function
      | Ast.Pipeline p -> (Fun.id, p)
      | Ast.Pipeline_Bang p -> (Exit.not, p)
    in

    let rec fold :
        Ast.and_or * ctx Exit.t -> Ast.pipeline Ast.and_or_list -> ctx Exit.t =
     fun (sep, exit_so_far) pipe ->
      match (sep, pipe) with
      | And, Nlist.Singleton (p, _) -> (
          match exit_so_far with
          | Exit.Zero ctx ->
              let f, p = pipeline p in
              f @@ execute_commands ctx sw p
          | v -> v)
      | Or, Nlist.Singleton (p, _) -> (
          match exit_so_far with
          | Exit.Zero _ as ctx -> ctx
          | _ ->
              let f, p = pipeline p in
              f @@ execute_commands ctx sw p)
      | Noand_or, Nlist.Singleton (p, _) ->
          let f, p = pipeline p in
          f @@ execute_commands ctx sw p
      | Noand_or, Nlist.Cons ((p, next_sep), rest) ->
          let f, p = pipeline p in
          fold (next_sep, f (execute_commands ctx sw p)) rest
      | And, Nlist.Cons ((p, next_sep), rest) -> (
          match exit_so_far with
          | Exit.Zero ctx ->
              let f, p = pipeline p in
              fold (next_sep, f (execute_commands ctx sw p)) rest
          | Exit.Nonzero _ as v -> v)
      | Or, Nlist.Cons ((p, next_sep), rest) -> (
          match exit_so_far with
          | Exit.Zero _ as exit_so_far -> fold (next_sep, exit_so_far) rest
          | Exit.Nonzero _ ->
              let f, p = pipeline p in
              fold (next_sep, f (execute_commands ctx sw p)) rest)
    in
    fold (Noand_or, Exit.zero ctx) c

  and handle_for_clause ctx = function
    | Ast.For_Name_DoGroup (_, (term, sep)) -> exec ctx (term, Some sep)
    | Ast.For_Name_In_WordList_DoGroup (Name name, wdlist, (term, sep)) ->
        let wdlist = Nlist.flatten @@ Nlist.map (word_glob_expand ctx) wdlist in
        Nlist.fold_left
          (fun _ word ->
            let s = S.update ctx.state ~param:name [ Ast.WordLiteral word ] in
            let ctx = { ctx with state = s } in
            exec ctx (term, Some sep))
          (Exit.zero ctx) wdlist

  and handle_if_clause ctx = function
    | Ast.If_then ((e1, sep1), (e2, sep2)) -> (
        let ctx = exec ctx (e1, Some sep1) in
        match ctx with
        | Exit.Zero ctx -> exec ctx (e2, Some sep2)
        | Exit.Nonzero { value = ctx; _ } -> Exit.zero ctx)
    | Ast.If_then_else ((e1, sep1), (e2, sep2), else_part) -> (
        let ctx = exec ctx (e1, Some sep1) in
        match ctx with
        | Exit.Zero ctx -> exec ctx (e2, Some sep2)
        | Exit.Nonzero { value = ctx; _ } -> handle_else_part ctx else_part)

  and handle_else_part ctx = function
    | Ast.Else (c, sep) -> exec ctx (c, Some sep)
    | Ast.Elif_then ((e1, sep1), (e2, sep2)) -> (
        let ctx = exec ctx (e1, Some sep1) in
        match ctx with
        | Exit.Zero ctx -> exec ctx (e2, Some sep2)
        | Exit.Nonzero { value = ctx; _ } -> Exit.zero ctx)
    | Ast.Elif_then_else ((e1, sep1), (e2, sep2), else_part) -> (
        let ctx = exec ctx (e1, Some sep1) in
        match ctx with
        | Exit.Zero ctx -> exec ctx (e2, Some sep2)
        | Exit.Nonzero { value = ctx; _ } -> handle_else_part ctx else_part)

  and handle_compound_command ctx = function
    | Ast.ForClause fc -> handle_for_clause ctx fc
    | Ast.IfClause if_ -> handle_if_clause ctx if_
    | _ as c ->
        Fmt.failwith "Compound command not supported: %a" yojson_pp
          (Ast.compound_command_to_yojson c)

  and needs_subshelling = function
    | [] -> false
    | Ast.WordSubshell _ :: _ -> true
    | Ast.WordDoubleQuoted word :: rest ->
        needs_subshelling word || needs_subshelling rest
    | Ast.WordSingleQuoted word :: rest ->
        needs_subshelling word || needs_subshelling rest
    | _ -> false

  and handle_subshell (ctx : ctx) wcs =
    let exec_subshell ~sw ctx s =
      let buf = Buffer.create 16 in
      let stdout = Eio.Flow.buffer_sink buf in
      let r, w = Eio_unix.pipe sw in
      Eio.Fiber.fork ~sw (fun () -> Eio.Flow.copy r stdout);
      let subshell_ctx = { ctx with stdout = Some w } in
      let _ = run (Exit.zero subshell_ctx) s in
      (ctx, Buffer.contents buf)
    in
    let rec run_subshells ~sw ran_subshell = function
      | [] -> []
      | Ast.WordSubshell s :: rest ->
          let _ctx, std = exec_subshell ~sw ctx s in
          ran_subshell := true;
          Ast.WordName (String.trim std) :: run_subshells ~sw ran_subshell rest
      | Ast.WordDoubleQuoted word :: rest ->
          let subshell_q = ref false in
          let res = run_subshells ~sw subshell_q word in
          if !subshell_q then res @ run_subshells ~sw subshell_q rest
          else Ast.WordDoubleQuoted res :: run_subshells ~sw subshell_q rest
      | Ast.WordSingleQuoted word :: rest ->
          let subshell_q = ref false in
          let res = run_subshells ~sw subshell_q word in
          if !subshell_q then res @ run_subshells ~sw subshell_q rest
          else Ast.WordSingleQuoted res :: run_subshells ~sw subshell_q rest
      | v :: rest -> v :: run_subshells ~sw ran_subshell rest
    in
    Eio.Switch.run @@ fun sw -> run_subshells ~sw (ref false) wcs

  and handle_word_components_to_string (ctx : ctx) wcs : string =
    if needs_subshelling wcs then begin
      let wcs = handle_subshell ctx wcs in
      Ast.word_components_to_string wcs
    end
    else Ast.word_components_to_string wcs

  and glob_expand ctx wc =
    handle_word_components_to_string ctx wc |> Globlon.glob |> Array.to_list

  and word_glob_expand (ctx : ctx) wc =
    if List.exists needs_glob_expansion wc then glob_expand ctx wc
    else [ handle_word_components_to_string ctx wc ]

  and collect_assignments ?(update = true) ctx =
    List.fold_left
      (fun ctx -> function
        | Ast.Prefix_assignment (Name param, v) ->
            (* Expand the values *)
            let ctx, v = expand_cst ctx v in
            let state =
              if update then S.update ctx.state ~param v else ctx.state
            in
            {
              ctx with
              state;
              local_state =
                (param, Ast.word_components_to_string v) :: ctx.local_state;
            }
        | _ -> ctx)
      ctx

  and args ctx swc =
    List.concat_map
      (function
        | Ast.Suffix_redirect _ -> []
        | Suffix_word wc ->
            let ctx, cst = expand_cst ctx wc in
            word_glob_expand ctx cst)
      swc

  and exec initial_ctx (ast : Ast.complete_command) =
    let command, _ = ast in
    let rec loop : Eio.Switch.t -> ctx -> Ast.clist -> ctx Exit.t =
     fun sw ctx -> function
       | Nlist.Singleton (c, _) -> handle_single_pipeline ~sw ctx c
       | Nlist.Cons ((c, (Semicolon | Nosep)), cs) -> (
           match handle_single_pipeline ~sw ctx c with
           | Exit.Zero ctx -> loop sw ctx cs
           | v -> v)
       | _ -> Fmt.failwith "Background tasks not implemented yet!"
    in
    Eio.Switch.run @@ fun sw -> loop sw initial_ctx command

  and execute ctx ast = exec ctx ast

  and run ctx ast =
    let ctx, cs =
      List.fold_left
        (fun (ctx, cs) command ->
          let ctx = Exit.value ctx in
          let exit = execute ctx command in
          match exit with
          | Exit.Nonzero { exit_code; message; should_exit; _ } -> (
              Option.iter (Fmt.epr "%s\n%!") message;
              match
                ( should_exit.interactive,
                  should_exit.non_interactive,
                  ctx.interactive )
              with
              | `Yes, _, true | _, `Yes, false -> Stdlib.exit exit_code
              | _ -> (exit, ast :: cs))
          | Exit.Zero _ as ctx -> (ctx, ast :: cs))
        (ctx, []) ast
    in
    (ctx, List.rev cs)
end