defmodule Diffuser.Generator.PromptRequestWorker do
  alias Diffuser.Generator
  alias Diffuser.Generator.PromptRequest
  alias DiffuserWeb.Endpoint
  alias Diffuser.Repo

  @path [:code.priv_dir(:diffuser), "python"] |> Path.join()

  def start(%PromptRequest{} = prompt_request) do
    with {:ok, active_prompt} <-
           update_and_broadcast_progress(prompt_request, "in_progress"),
         {:ok, _file_location} <- call_python(:test_script, :test_func, active_prompt),
         %PromptRequest{} = prompt_request_with_results <-
           write_and_save_images(active_prompt),
         {:ok, completed_prompt} <-
           update_and_broadcast_progress(prompt_request_with_results, "finished") do
      {:ok, completed_prompt |> Repo.preload(:images)}
    else
      nil ->
        raise("prompt not found")

      {:error, message} ->
        raise(message)
    end
  end

  defp update_and_broadcast_progress(prompt_request, "in_progress"),
    do:
      update_and_broadcast_progress(prompt_request, %{
        status: "in_progress",
        began_at: NaiveDateTime.utc_now()
      })

  defp update_and_broadcast_progress(prompt_request, "finished"),
    do:
      update_and_broadcast_progress(prompt_request, %{
        status: "finished",
        ended_at: NaiveDateTime.utc_now()
      })

  defp update_and_broadcast_progress(%PromptRequest{id: id} = prompt_request, attrs) do
    {:ok, new_prompt} = Generator.update_prompt_request(prompt_request, attrs)

    :ok = Endpoint.broadcast("request:#{id}", "request", %{prompt_request: new_prompt})

    {:ok, new_prompt}
  end

  defp call_python(_module, _func, %PromptRequest{
         id: prompt_id,
         prompt: prompt,
         steps: steps,
         guidance_scale: guidance_scale
       }) do
    port =
      Port.open(
        {:spawn,
         ~s(python #{@path}/stable_diffusion.py --prompt "#{prompt}" --output "#{@path}/#{prompt_id}.png" --num-inference-steps #{steps} --guidance-scale #{guidance_scale})},
        [:binary]
      )

    python_loop(port, prompt_id)
  end

  defp python_loop(port, prompt_id) do
    receive do
      {^port, {:data, ":finished" <> msg}} ->
        {:ok, msg}

      {^port, {:data, ":step" <> step}} ->
        {:ok, prompt_request} =
          Generator.update_prompt_request(prompt_id, %{completed_steps: step})

        Endpoint.broadcast(
          "request:#{prompt_id}",
          "progress",
          %{prompt_request: prompt_request |> Repo.preload(:images)}
        )

        python_loop(port, prompt_id)

      {^port, _result} ->
        python_loop(port, prompt_id)
    end
  end

  defp write_and_save_images(%PromptRequest{id: id}) do
    file_path = "#{@path}/#{id}.png"

    with {:ok, body} <- File.read(file_path),
         {:ok, _result} <-
           Generator.create_prompt_request_result(
             id,
             %{
               file_name: "#{id}.png",
               filename: "#{id}.png",
               binary: body
             }
           ),
         :ok <- File.rm(file_path) do
      Generator.get_prompt_request!(id)
    end
  end
end