#!/usr/bin/env ruby
# frozen_string_literal: true

require "landlock"
require "msgpack"
require "socket"
require "tmpdir"
require "vips"

raise LoadError, "libvips 8.13 or newer is required" if !Vips.at_least_libvips?(8, 13)

class DiscourseVipsWorker
  class InvalidImage < StandardError
  end

  LINUX_PLATFORM = RUBY_PLATFORM.include?("linux")
  private_constant :LINUX_PLATFORM

  class Operations
    DEFAULT_READ_PATHS = %w[/bin /lib /lib64 /usr].freeze
    private_constant :DEFAULT_READ_PATHS

    # Dynamically loaded libvips modules use this cache to locate their shared libraries.
    DYNAMIC_LINKER_CACHE_PATH = "/etc/ld.so.cache"
    private_constant :DYNAMIC_LINKER_CACHE_PATH

    WARM_OPERATIONS = %w[text gravity flatten pngsave thumbnail].freeze
    private_constant :WARM_OPERATIONS

    def self.warm
      WARM_OPERATIONS.each do |name|
        operation = Vips.vips_operation_new(name)
        raise "unable to initialize libvips operation #{name}" if operation.null?

        GObject.g_object_unref(operation)
      end
    end

    def self.build(command)
      operation, *arguments = command

      case operation
      when "version"
        version
      when "letter-avatar"
        letter_avatar(*arguments)
      when "dominant-color"
        dominant_color(*arguments)
      else
        raise ArgumentError, "unsupported libvips operation"
      end
    end

    def initialize(read_paths: [], write_paths: [], &action)
      @read_paths = read_paths
      @write_paths = write_paths
      @action = action
    end

    def call
      @action.call
    end

    def sandbox_paths(scratch:)
      {
        "read" => existing_paths([*DEFAULT_READ_PATHS, DYNAMIC_LINKER_CACHE_PATH, *@read_paths]),
        "write" => existing_paths([scratch, *@write_paths]),
      }
    end

    def self.version
      new { Vips.version_string }
    end
    private_class_method :version

    def self.letter_avatar(letter, output_path, background_color, font, font_path)
      new(
        read_paths: ["/etc/fonts", "/var/cache/fontconfig", font_path],
        write_paths: [File.dirname(output_path)],
      ) do
        if !letter || !output_path || !font_path
          raise ArgumentError, "invalid letter avatar arguments"
        end
        raise ArgumentError, "font file does not exist" if !File.file?(font_path)
        raise ArgumentError, "invalid background color" if !background_color.match?(/\A\h{6}\z/)

        background = background_color.scan(/\h{2}/).map { |channel| Integer(channel, 16) }

        block_loaders
        markup = %(<span foreground="#ffffff" alpha="80%">#{letter}</span>)
        text = Vips::Image.text(markup, font:, dpi: 72, fontfile: font_path, rgba: true)
        canvas =
          text.gravity(:centre, 360, 360, extend: :background, background: [*background, 255])
        canvas.flatten(background:).pngsave(output_path, compression: 6)
        nil
      end
    end
    private_class_method :letter_avatar

    def self.dominant_color(input_path)
      new(read_paths: [input_path]) do
        raise ArgumentError, "invalid dominant color arguments" if !input_path

        block_loaders(
          allowed: %w[
            VipsForeignLoadJpeg
            VipsForeignLoadPng
            VipsForeignLoadNsgif
            VipsForeignLoadWebp
            VipsForeignLoadHeif
            VipsForeignLoadJxl
          ],
        )
        components =
          begin
            image = Vips::Image.new_from_file(input_path, access: :sequential)
            if image.has_alpha?
              alpha_weighted_components(image)
            elsif %i[float double].include?(image.format)
              normalize_color(image)
                .thumbnail_image(1, height: 1, size: :force)
                .getpoint(0, 0)
                .first(3)
                .map(&:round)
            else
              thumbnail = Vips::Image.thumbnail(input_path, 1, height: 1, size: :force)
              normalize_color(thumbnail).getpoint(0, 0).first(3).map(&:round)
            end
          rescue Vips::Error => error
            raise InvalidImage, error.message
          end
        format("%02X%02X%02X", *components)
      end
    end
    private_class_method :dominant_color

    def self.normalize_color(image)
      image = image.colourspace(:srgb)
      image = (image * 255).round(:rint) if %i[float double].include?(image.format)
      image.cast(:uchar, shift: true)
    end
    private_class_method :normalize_color

    def self.alpha_weighted_components(image)
      alpha = image[image.bands - 1]
      colors = normalize_color(image.extract_band(0, n: image.bands - 1))
      weighted_colors = colors * alpha
      statistics = alpha.bandjoin(weighted_colors).stats
      alpha_sum = statistics.getpoint(2, 1).first
      return 0, 0, 0 if alpha_sum.zero?

      3.times.map { |band| (statistics.getpoint(2, band + 2).first / alpha_sum).round }
    end
    private_class_method :alpha_weighted_components

    def self.block_loaders(allowed: [])
      Vips.block_untrusted(true)
      Vips.block("VipsForeignLoad", true)
      allowed.each { |loader| Vips.block(loader, false) }
    end
    private_class_method :block_loaders

    private

    def existing_paths(paths)
      paths
        .filter { |path| path.to_s != "" && File.exist?(path) }
        .map { |path| File.expand_path(path) }
        .uniq
    end
  end

  RLIMITS = {
    cpu_seconds: 300,
    memory_bytes: 4 * 1024 * 1024 * 1024,
    file_size_bytes: 10 * 1024 * 1024 * 1024,
    open_files: 1024,
  }.freeze
  private_constant :RLIMITS

  def self.run(server_fd:, owner_fd:)
    parent_pid = Process.ppid
    Landlock::Native.set_parent_death_signal! if LINUX_PLATFORM
    exit! 1 if Process.ppid != parent_pid

    Process.setproctitle("discourse vips worker")
    new(server: UNIXServer.for_fd(server_fd), owner_reader: IO.for_fd(owner_fd)).run
  end

  def initialize(server:, owner_reader:)
    @server = server
    @socket_path = server.local_address.unix_path
    @socket_identity = File.stat(@socket_path).then { |stat| [stat.dev, stat.ino] }
    @owner_reader = owner_reader
  end

  def run
    Signal.trap("CHLD") { reap_connection_processes }
    Operations.warm

    loop do
      ready_ios = IO.select([@server, @owner_reader]).first
      break if ready_ios.include?(@owner_reader)

      connection = @server.accept
      fork_connection_process(connection)
    end
  ensure
    @server.close unless @server.closed?
    @owner_reader.close unless @owner_reader.closed?
    remove_socket_directory
    Process.kill("KILL", -Process.pid)
  end

  private

  def fork_connection_process(connection)
    parent_pid = Process.pid
    fork do
      Landlock::Native.set_parent_death_signal! if LINUX_PLATFORM
      exit! 1 if Process.ppid != parent_pid
      Signal.trap("CHLD", "DEFAULT")
      @server.close
      @owner_reader.close
      handle_connection(connection)
      exit! 0
    end

    connection.close
  end

  def reap_connection_processes
    while Process.waitpid(-1, Process::WNOHANG)
    end
  rescue Errno::ECHILD
  end

  def handle_connection(connection)
    request = MessagePack.unpack(connection.read)
    operation = Operations.build(request.fetch("command"))
    response = run_sandboxed(request:, operation:)
    write_response(connection, response)
  rescue StandardError => error
    write_response(connection, "status" => "error", "message" => error.message.byteslice(0, 4096))
  ensure
    connection.close unless connection.closed?
  end

  def write_response(connection, response)
    connection.write(MessagePack.pack(response))
  rescue IOError, SystemCallError
  end

  def run_sandboxed(request:, operation:)
    Dir.mktmpdir("discourse-vips-") do |scratch|
      sandbox_paths = operation.sandbox_paths(scratch:)
      result =
        Landlock.fork(
          read: sandbox_paths.fetch("read"),
          write: sandbox_paths.fetch("write"),
          execute: [],
          timeout: Float(request.fetch("timeout")),
          env: child_environment(scratch),
          unsetenv_others: true,
          rlimits: RLIMITS,
          seccomp_deny_network: LINUX_PLATFORM,
          on_unsupported: :run_without_landlock,
        ) do |stdout, _stderr|
          apply_nice(request)
          stdout.write(MessagePack.pack(operation_response(operation)))
        end

      if result.timed_out?
        { "status" => "timeout" }
      elsif result.success?
        MessagePack.unpack(result.stdout)
      else
        { "status" => "error", "message" => "libvips operation failed" }
      end
    end
  end

  def remove_socket_directory
    stat = File.stat(@socket_path)
    return if [stat.dev, stat.ino] != @socket_identity

    File.unlink(@socket_path)
    Dir.rmdir(File.dirname(@socket_path))
  rescue SystemCallError
  end

  def child_environment(scratch)
    {
      "HOME" => scratch,
      "TMPDIR" => scratch,
      "XDG_CACHE_HOME" => scratch,
      "MALLOC_ARENA_MAX" => "2",
    }
  end

  def apply_nice(request)
    Process.setpriority(Process::PRIO_PROCESS, 0, Integer(request.fetch("nice"))) if request["nice"]
  end

  def operation_response(operation)
    value = operation.call
    { "status" => "ok", "value" => value }
  rescue InvalidImage => error
    { "status" => "invalid_image", "message" => error.message.byteslice(0, 4096) }
  rescue StandardError => error
    { "status" => "error", "message" => error.message.byteslice(0, 4096) }
  end
end

if $PROGRAM_NAME == __FILE__
  server_fd, owner_fd = ARGV.map { |argument| Integer(argument) }
  DiscourseVipsWorker.run(server_fd:, owner_fd:)
end