defmodule Nola.Plugins.Link.Image do
  require Logger
  @behaviour Nola.Plugins.Link

  @impl true
  def match(_, _), do: false

  @impl true
  def post_match(_url, "image/" <> _, _header, _opts) do
    {:body, nil}
  end

  def post_match(_, _, _, _), do: false

  @impl true
  def post_expand(_url, bytes, _, opts) do
    pil_process = Keyword.get(opts, :pil_process, {:pil, :"py@127.0.0.1"})
    clip_ask_process = Keyword.get(opts, :clip_ask_process, {:clip_ask, :"py@127.0.0.1"})

    img2txt_process =
      Keyword.get(opts, :img2txt_process, {:image_to_text_vit_gpt2, :"py@127.0.0.1"})

    tasks =
      [
        Task.async(fn ->
          {:ok, pil} = GenServer.call(pil_process, {:run, bytes})

          pil =
            pil
            |> Enum.map(fn {k, v} -> {String.to_atom(to_string(k)), v} end)

          pil
        end),
        Task.async(fn ->
          {:ok, descr} = GenServer.call(img2txt_process, {:run, bytes})
          {:img2txt, to_string(descr)}
        end),
        Task.async(fn ->
          {:ok, prompts} = GenServer.call(clip_ask_process, {:run, bytes})

          prompts =
            prompts
            |> Enum.sort_by(&elem(&1, 1), &>=/2)
            |> Enum.take(3)
            |> Enum.map(&to_string(elem(&1, 0)))
            |> Enum.join(", ")

          {:prompts, prompts}
        end)
      ]
      |> Task.yield_many(5000)
      |> Enum.map(fn {task, res} ->
        res || Task.shutdown(task, :brutal_kill)
      end)

    results = Enum.into(List.flatten(for({:ok, value} <- tasks, do: value)), Map.new())
    img2txt = Map.get(results, :img2txt)
    prompts = Map.get(results, :prompts)

    pil =
      if Map.get(results, :width) do
        animated = if Map.get(results, :animated), do: " animated", else: ""
        "#{Map.get(results, :width, 0)}x#{Map.get(results, :height, 0)}#{animated} — "
      else
        ""
      end

    descr =
      cond do
        img2txt && prompts ->
          "#{pil}#{prompts} — #{img2txt}"

        img2txt ->
          "#{pil}#{img2txt}"

        prompts ->
          "#{pil}#{prompts}"

        pil != "" ->
          "#{pil}"

        true ->
          nil
      end

    if descr do
      {:ok, "image: #{descr}"}
    else
      :error
    end
  end
end