Websites/app/logical/iqdb.rb

# frozen_string_literal: true

module Iqdb
  class Error < StandardError; end

  module_function

  IQDB_NUM_PIXELS = 128

  def endpoint
    Websites.config.iqdb_server
  end

  def enabled?
    endpoint.present?
  end

  def make_request(path, request_type, body = nil)
    conn = Faraday.new(Websites.config.faraday_options)
    conn.send(request_type, endpoint + path, body&.to_json, { content_type: "application/json" })
  rescue Faraday::Error
    raise(Error, "This service is temporarily unavailable. Please try again later.")
  end

  def update_image(image)
    Websites.config.yiffy2_storage.get(image.path) do |file|
      Tempfile.open("yiffy2-#{image.md5}") do |tempfile|
        file.binmode
        tempfile.binmode
        tempfile.write(file.read)
        tempfile.rewind
        thumb = generate_thumbnail(tempfile.path)
        raise(Error, "failed to generate thumb for #{image.iqdb_id}") unless thumb
        response = make_request("/images/#{image.iqdb_id}", :post, get_channels_data(thumb))
        raise(Error, "iqdb request failed") if response.status != 200
      end
    end
  end

  def remove_image(id)
    response = make_request("/images/#{id}", :delete)
    raise(Error, "iqdb request failed") if response.status != 200
  end

  def query_url(url, score_cutoff)
    file = FileDownload.new(url).download!
    query_file(file, score_cutoff)
  end

  def query_image(image, score_cutoff)
    Websites.config.yiffy2_storage.get(image.path) do |file|
      query_file(file, score_cutoff)
    end
  end

  def query_file(file, score_cutoff)
    Tempfile.open("yiffy2-#{file.md5}") do |tempfile|
      file.binmode
      tempfile.binmode
      tempfile.write(file.read)
      tempfile.rewind
      thumb = generate_thumbnail(tempfile.path)
      raise(Error, "failed to generate thumb for #{file.md5}") unless thumb

      response = make_request("/query", :post, get_channels_data(thumb))
      raise(Error, "iqdb request failed") if response.status != 200

      process_iqdb_result(JSON.parse(response.body), score_cutoff)
    end
  end

  def query_hash(hash, score_cutoff)
    response = make_request("/query", :post, { hash: hash })
    raise(Error, "iqdb request failed") if response.status != 200

    process_iqdb_result(JSON.parse(response.body), score_cutoff)
  end

  def process_iqdb_result(json, score_cutoff)
    raise(Error, "Server returned an error. Most likely the url is not found.") unless json.is_a?(Array)

    json.filter! { |entry| (entry["score"] || 0) >= (score_cutoff.presence || 50).to_i }
    json.map do |x|
      x["image"] = APIImage.find_by!(iqdb_id: x["post_id"])
      x
    rescue ActiveRecord::RecordNotFound
      nil
    end.compact
  end

  def generate_thumbnail(file_path)
    Vips::Image.thumbnail(file_path, IQDB_NUM_PIXELS, height: IQDB_NUM_PIXELS, size: :force)
  rescue Vips::Error => e
    Rails.logger.error("failed to generate thumbnail for #{file_path}")
    Rails.logger.error(e)
    nil
  end

  def get_channels_data(thumbnail)
    r = []
    g = []
    b = []
    is_grayscale = thumbnail.bands == 1
    thumbnail.to_a.each do |data|
      data.each do |rgb|
        r << rgb[0]
        g << (is_grayscale ? rgb[0] : rgb[1])
        b << (is_grayscale ? rgb[0] : rgb[2])
      end
    end
    { channels: { r: r, g: g, b: b } }
  end
end