require "c/processthreadsapi"
require "c/handleapi"
require "c/jobapi2"
require "c/synchapi"
require "c/tlhelp32"
require "c/debugapi"
require "process/shell"
require "crystal/atomic_semaphore"

struct Crystal::System::Process
  {% if host_flag?(:windows) %}
    HOST_PATH_DELIMITER = ';'
  {% else %}
    HOST_PATH_DELIMITER = ':'
  {% end %}

  getter pid : LibC::DWORD
  @thread_id : LibC::DWORD
  @process_handle : LibC::HANDLE
  @job_object : LibC::HANDLE
  @completion_key = IOCP::CompletionKey.new(:process_run)

  @@interrupt_handler : Proc(::Process::ExitReason, Nil)?
  @@interrupt_count = Crystal::AtomicSemaphore.new
  @@win32_interrupt_handler : LibC::PHANDLER_ROUTINE?
  @@setup_interrupt_handler = Atomic(Bool).new(false)
  @@last_interrupt = ::Process::ExitReason::Interrupted

  def initialize(process_info)
    @pid = process_info.dwProcessId
    @thread_id = process_info.dwThreadId
    @process_handle = process_info.hProcess

    @job_object = LibC.CreateJobObjectW(nil, nil)

    # enable IOCP notifications
    config_job_object(
      LibC::JOBOBJECTINFOCLASS::AssociateCompletionPortInformation,
      LibC::JOBOBJECT_ASSOCIATE_COMPLETION_PORT.new(
        completionKey: @completion_key.as(Void*),
        completionPort: Crystal::EventLoop.current.iocp_handle,
      ),
    )

    # but not for any child processes
    config_job_object(
      LibC::JOBOBJECTINFOCLASS::ExtendedLimitInformation,
      LibC::JOBOBJECT_EXTENDED_LIMIT_INFORMATION.new(
        basicLimitInformation: LibC::JOBOBJECT_BASIC_LIMIT_INFORMATION.new(
          limitFlags: LibC::JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK,
        ),
      ),
    )

    if LibC.AssignProcessToJobObject(@job_object, @process_handle) == 0
      raise RuntimeError.from_winerror("AssignProcessToJobObject")
    end

    if LibC.ResumeThread(process_info.hThread) == 0xFFFFFFFF_u32
      raise RuntimeError.from_winerror("ResumeThread")
    end

    close_handle(process_info.hThread)
  end

  private def config_job_object(kind, info)
    if LibC.SetInformationJobObject(@job_object, kind, pointerof(info), sizeof(typeof(info))) == 0
      raise RuntimeError.from_winerror("SetInformationJobObject")
    end
  end

  def release
    return if @process_handle == LibC::HANDLE.null
    close_handle(@process_handle)
    @process_handle = LibC::HANDLE.null
    close_handle(@job_object)
    @job_object = LibC::HANDLE.null
  end

  def wait
    @completion_key.fiber = ::Fiber.current

    if LibC.GetExitCodeProcess(@process_handle, out exit_code) == 0
      raise RuntimeError.from_winerror("GetExitCodeProcess")
    end

    unless exit_code == LibC::STILL_ACTIVE
      # MT race? another thread received the completion event and will resume
      # the fiber, we must suspend the current fiber before we can return
      ::Fiber.suspend unless @completion_key.reset_fiber?

      return exit_code
    end

    # let `@job_object` do its job
    # TODO: message delivery is "not guaranteed"; does it ever happen? Are we
    # stuck forever in that case?
    # (https://learn.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-jobobject_associate_completion_port)
    ::Fiber.suspend

    # If the IOCP notification is delivered before the process fully exits,
    # wait for it
    if LibC.WaitForSingleObject(@process_handle, LibC::INFINITE) != LibC::WAIT_OBJECT_0
      raise RuntimeError.from_winerror("WaitForSingleObject")
    end

    # WaitForSingleObject returns immediately once ExitProcess is called in the child, but
    # the process still has yet to be destructed by the OS and have it's memory unmapped.
    # Since the semantics on unix are that the resources of a process have been released once
    # waitpid returns, we wait 5 milliseconds to attempt to replicate this behaviour.
    sleep 5.milliseconds

    if LibC.GetExitCodeProcess(@process_handle, pointerof(exit_code)) == 0
      raise RuntimeError.from_winerror("GetExitCodeProcess")
    end
    if exit_code == LibC::STILL_ACTIVE
      raise "BUG: Process still active"
    end
    exit_code
  end

  def exists?
    Crystal::System::Process.exists?(@pid)
  end

  def terminate(*, graceful)
    LibC.TerminateProcess(@process_handle, 1)
  end

  def self.exit(status : Int32)
    LibC.exit(status)
  end

  def self.exit(status : ::Process::Status)
    exit status.system_exit_status.to_i32!
  end

  def self.pid
    LibC.GetCurrentProcessId
  end

  def self.pgid
    raise NotImplementedError.new("Process.pgid")
  end

  def self.pgid(pid)
    raise NotImplementedError.new("Process.pgid")
  end

  def self.ppid
    pid = self.pid
    each_process_entry do |pe|
      return pe.th32ParentProcessID if pe.th32ProcessID == pid
    end
    raise RuntimeError.new("Cannot locate current process")
  end

  private def self.each_process_entry(&)
    h = LibC.CreateToolhelp32Snapshot(LibC::TH32CS_SNAPPROCESS, 0)
    raise RuntimeError.from_winerror("CreateToolhelp32Snapshot") if h == LibC::INVALID_HANDLE_VALUE

    begin
      pe = LibC::PROCESSENTRY32W.new(dwSize: sizeof(LibC::PROCESSENTRY32W))
      if LibC.Process32FirstW(h, pointerof(pe)) != 0
        while true
          yield pe
          break if LibC.Process32NextW(h, pointerof(pe)) == 0
        end
      end
    ensure
      LibC.CloseHandle(h)
    end
  end

  def self.signal(pid, signal)
    raise NotImplementedError.new("Process.signal")
  end

  @[Deprecated("Use `#on_terminate` instead")]
  def self.on_interrupt(&handler : ->) : Nil
    on_terminate do |reason|
      handler.call if reason.interrupted?
    end
  end

  def self.on_terminate(&@@interrupt_handler : ::Process::ExitReason ->) : Nil
    restore_interrupts!
    @@win32_interrupt_handler = handler = LibC::PHANDLER_ROUTINE.new do |event_type|
      @@last_interrupt = case event_type
                         when LibC::CTRL_C_EVENT, LibC::CTRL_BREAK_EVENT
                           ::Process::ExitReason::Interrupted
                         when LibC::CTRL_CLOSE_EVENT
                           ::Process::ExitReason::TerminalDisconnected
                         when LibC::CTRL_LOGOFF_EVENT, LibC::CTRL_SHUTDOWN_EVENT
                           ::Process::ExitReason::SessionEnded
                         else
                           next 0
                         end
      @@interrupt_count.signal
      1
    end
    LibC.SetConsoleCtrlHandler(handler, 1)
  end

  def self.ignore_interrupts! : Nil
    remove_interrupt_handler
    LibC.SetConsoleCtrlHandler(nil, 1)
  end

  def self.restore_interrupts! : Nil
    remove_interrupt_handler
    LibC.SetConsoleCtrlHandler(nil, 0)
  end

  private def self.remove_interrupt_handler
    if old = @@win32_interrupt_handler
      LibC.SetConsoleCtrlHandler(old, 0)
      @@win32_interrupt_handler = nil
    end
  end

  def self.start_interrupt_loop : Nil
    return if @@setup_interrupt_handler.swap(true, :relaxed)

    spawn(name: "interrupt-signal-loop") do
      while true
        @@interrupt_count.wait { sleep 50.milliseconds }

        if handler = @@interrupt_handler
          non_nil_handler = handler # if handler is closured it will also have the Nil type
          int_type = @@last_interrupt
          spawn do
            non_nil_handler.call int_type
          rescue ex
            ex.inspect_with_backtrace(STDERR)
            STDERR.puts("FATAL: uncaught exception while processing interrupt handler, exiting")
            STDERR.flush
            LibC._exit(1)
          end
        end
      end
    end
  end

  def self.debugger_present? : Bool
    LibC.IsDebuggerPresent != 0
  end

  def self.exists?(pid)
    handle = LibC.OpenProcess(LibC::PROCESS_QUERY_INFORMATION, 0, pid)
    return false unless handle
    begin
      if LibC.GetExitCodeProcess(handle, out exit_code) == 0
        raise RuntimeError.from_winerror("GetExitCodeProcess")
      end
      exit_code == LibC::STILL_ACTIVE
    ensure
      close_handle(handle)
    end
  end

  def self.times
    if LibC.GetProcessTimes(LibC.GetCurrentProcess, out create, out exit, out kernel, out user) == 0
      raise RuntimeError.from_winerror("GetProcessTimes")
    end
    ::Process::Tms.new(
      Crystal::System::Time.filetime_to_f64secs(user),
      Crystal::System::Time.filetime_to_f64secs(kernel),
      0,
      0)
  end

  def self.fork
    raise NotImplementedError.new("Process.fork")
  end

  def self.fork(&)
    raise NotImplementedError.new("Process.fork")
  end

  private def self.handle_from_io(io : IO::FileDescriptor, parent_io)
    source_handle =
      if io.is_a?(File) && !io.system_blocking? && !io.closed?
        dup_handle = reopen_file_as_blocking(io, parent_io == STDIN, "Process.run")
      else
        io.windows_handle
      end

    cur_proc = LibC.GetCurrentProcess
    if LibC.DuplicateHandle(cur_proc, source_handle, cur_proc, out new_handle, 0, true, LibC::DUPLICATE_SAME_ACCESS) == 0
      raise RuntimeError.from_winerror("DuplicateHandle")
    end

    {new_handle, dup_handle}
  end

  def self.spawn(prepared_args, shell, env, clear_env, input, output, error, chdir, &)
    startup_info = LibC::STARTUPINFOW.new
    startup_info.cb = sizeof(LibC::STARTUPINFOW)
    startup_info.dwFlags = LibC::STARTF_USESTDHANDLES

    startup_info.hStdInput, dup_input = handle_from_io(input, STDIN)
    startup_info.hStdOutput, dup_output = handle_from_io(output, STDOUT)
    startup_info.hStdError, dup_error = handle_from_io(error, STDERR)

    process_info = LibC::PROCESS_INFORMATION.new

    prepared_args = ::Process.quote_windows(prepared_args) unless prepared_args.is_a?(String)

    if LibC.CreateProcessW(
         nil, System.to_wstr(prepared_args), nil, nil, true, LibC::CREATE_SUSPENDED | LibC::CREATE_UNICODE_ENVIRONMENT,
         Env.make_env_block(env, clear_env), chdir.try { |str| System.to_wstr(str) } || Pointer(UInt16).null,
         pointerof(startup_info), pointerof(process_info)
       ) == 0
      error = WinError.value
      if ::File::NotFoundError.os_error?(error) || ::File::AccessDeniedError.os_error?(error) || error.in?(WinError::ERROR_BAD_EXE_FORMAT, WinError::ERROR_INVALID_PARAMETER)
        yield error, prepared_args
      else
        raise IO::Error.from_os_error("Error executing process: '#{prepared_args}'", error)
      end
    end

    close_handle(startup_info.hStdInput)
    close_handle(startup_info.hStdOutput)
    close_handle(startup_info.hStdError)

    close_handle(dup_input) if dup_input
    close_handle(dup_output) if dup_output
    close_handle(dup_error) if dup_error

    process_info
  end

  def self.prepare_args(command : String, args : Enumerable(String)?, shell : Bool)
    if shell
      if args
        raise NotImplementedError.new("Process with args and shell: true is not supported on Windows")
      end
      command
    else
      command_args = [command]
      command_args.concat(args) if args
      prepare_args(command_args)
    end
  end

  def self.prepare_args(args : Enumerable(String))
    # Disable implicit execution of batch files (https://github.com/crystal-lang/crystal/issues/14536)
    #
    # > `CreateProcessW()` implicitly spawns `cmd.exe` when executing batch files (`.bat`, `.cmd`, etc.), even if the application didn’t specify them in the command line.
    # > The problem is that the `cmd.exe` has complicated parsing rules for the command arguments, and programming language runtimes fail to escape the command arguments properly.
    # > Because of this, it’s possible to inject commands if someone can control the part of command arguments of the batch file.
    # https://flatt.tech/research/posts/batbadbut-you-cant-securely-execute-commands-on-windows/
    command = args.first
    if command.rstrip(". ").byte_slice?(-4, 4).try(&.downcase).in?(".bat", ".cmd")
      raise ::File::Error.from_os_error("Error executing process", WinError::ERROR_BAD_EXE_FORMAT, file: command)
    end

    args
  end

  private def self.try_replace(command, prepared_args, env, clear_env, input, output, error, chdir)
    old_input_fd = reopen_io(input, ORIGINAL_STDIN)
    old_output_fd = reopen_io(output, ORIGINAL_STDOUT)
    old_error_fd = reopen_io(error, ORIGINAL_STDERR)

    ENV.clear if clear_env
    env.try &.each do |key, val|
      if val
        ENV[key] = val
      else
        ENV.delete key
      end
    end

    if prepared_args.is_a?(String)
      command = System.to_wstr(prepared_args)
      argv = [command]
    else
      command = System.to_wstr(prepared_args[0])
      argv = prepared_args.map { |arg| System.to_wstr(arg) }
    end

    argv << Pointer(LibC::WCHAR).null

    if chdir
      ::Dir.cd(chdir) do
        LibC._wexecvp(command, argv)
      end
    else
      LibC._wexecvp(command, argv)
    end

    # exec failed; restore the original C runtime file descriptors
    errno = Errno.value
    LibC._dup2(old_input_fd, 0)
    LibC._dup2(old_output_fd, 1)
    LibC._dup2(old_error_fd, 2)
    errno
  end

  def self.replace(command, args, shell, env, clear_env, input, output, error, chdir) : NoReturn
    prepared_args = prepare_args(command, args, shell)
    errno = try_replace(command, prepared_args, env, clear_env, input, output, error, chdir)
    raise_exception_from_errno(command, errno)
  end

  private def self.raise_exception_from_errno(command, errno = Errno.value)
    if ::File::NotFoundError.os_error?(errno) || ::File::AccessDeniedError.os_error?(errno)
      raise ::File::Error.from_os_error("Error executing process", errno, file: command)
    else
      raise IO::Error.from_os_error("Error executing process: '#{command}'", errno)
    end
  end

  # Replaces the C standard streams' file descriptors, not Win32's, since
  # `try_replace` uses the C `LibC._wexecvp` and only cares about the former.
  # Returns a duplicate of the original file descriptor
  private def self.reopen_io(src_io : IO::FileDescriptor, dst_io : IO::FileDescriptor)
    dst_fd =
      case dst_io
      when ORIGINAL_STDIN  then 0
      when ORIGINAL_STDOUT then 1
      when ORIGINAL_STDERR then 2
      else
        raise "BUG: Invalid destination IO"
      end

    src_fd =
      case src_io
      when STDIN  then 0
      when STDOUT then 1
      when STDERR then 2
      else
        handle =
          case src_io
          when .system_blocking?, .closed?
            src_io.windows_handle
          when File
            reopen_file_as_blocking(src_io, dst_fd == 0, "Process.exec")
          else
            raise IO::Error.new("Non-blocking streams are not supported in `Process.exec`", target: src_io)
          end
        LibC._open_osfhandle(handle, 0)
      end

    return src_fd if dst_fd == src_fd

    orig_src_fd = LibC._dup(src_fd)

    if LibC._dup2(src_fd, dst_fd) == -1
      raise IO::Error.from_errno("Failed to replace C file descriptor", target: dst_io)
    end

    orig_src_fd
  end

  # The arguments for input, output and error must be handles without
  # FILE_FLAG_OVERLAPPED.
  private def self.reopen_file_as_blocking(file, for_stdin, method_name)
    access = for_stdin ? LibC::FILE_GENERIC_READ : LibC::FILE_GENERIC_WRITE
    handle = LibC.ReOpenFile(file.windows_handle, access, LibC::DEFAULT_SHARE_MODE, 0)
    if handle == LibC::INVALID_HANDLE_VALUE
      raise IO::Error.from_winerror("Failed to reopen non-blocking File as blocking", target: file)
    end
    handle
  end

  def self.chroot(path)
    raise NotImplementedError.new("Process.chroot")
  end
end

private def close_handle(handle)
  if LibC.CloseHandle(handle) == 0
    raise RuntimeError.from_winerror("CloseHandle")
  end
end
