defmodule MedicalTranscriptionWeb.Components.TranscriptionTextComponent do
  @moduledoc """
  Represents a portion of transcribed text and its codes and starts a task to determine keywords within the text.
  """

  use MedicalTranscriptionWeb, :live_component
  import MedicalTranscriptionWeb.Components
  import MedicalTranscriptionWeb.Components.KeywordHighlighter
  alias AudioTagger.Structs.TagResult

  @impl Phoenix.LiveComponent
  def update(assigns, socket) do
    self_pid = self()

    socket =
      if Map.has_key?(assigns, :row) do
        socket
        |> assign(:row, assigns.row)
        |> assign_async(:keywords, fn -> find_keywords(self_pid, assigns.row.text) end)
      else
        socket
      end

    {:ok, socket}
  end

  @impl Phoenix.LiveComponent
  def render(assigns) do
    ~H"""
    <div class="flex gap-12 pb-10 border-b border-[#444444]/20">
      <div class="flex-1 flex flex-col gap-4">
        <p class="text-[32px] leading-normal font-semibold">
          <%= if !@row.is_summary do %>
            <%= @row.start_mark %> - <%= @row.end_mark %>
          <% end %>
        </p>
        <p class="text-[28px] leading-normal text-type-black-tertiary">
          <.async_result :let={keywords} assign={@keywords}>
            <:loading><%= @row.text %></:loading>
            <.highlight text={@row.text} keywords={keywords} />
          </.async_result>
        </p>

        <div class="flex items-center gap-2 text-sm italic">
          <.async_result assign={@keywords}>
            <:loading>
              <img src={~p"/images/loading.svg"} width="16" class="animate-spin" />
              Finding keywords...
            </:loading>
            <:failed :let={reason}>There was an error finding keywords: <%= reason %></:failed>
          </.async_result>
        </div>
      </div>

      <div class="flex-1 flex flex-col items-stretch gap-3">
        <%= for %TagResult{code: code, label: label, score: score} <- @row.tags do %>
          <.tag_result code={code} label={label} score={score} text={@row.text} />
        <% end %>
      </div>
    </div>
    """
  end

  defp find_keywords(live_view_pid, text) do
    # First, we use token classification to determine parts of speech and then retrieve the verb and adjective+noun phrases.
    %{entities: entities} = Nx.Serving.batched_run(TokenClassificationServing, text)
    phrases = AudioTagger.KeywordFinder.cleanup_phrases(entities)

    # Then, we use one of two processes to determine which to show as keywords:
    # 1. A slower process that looks to classify the text by the extracted phrases.
    # serving = AudioTagger.KeywordFinder.prepare_zero_shot_classification_serving(phrases)
    # %{predictions: predictions} = Nx.Serving.run(serving, text)

    # 2. A fast process finding the phrase closest in vector space to the whole text.
    predictions = AudioTagger.KeywordFinder.find_most_similar_label(text, phrases, 2)

    # For now, retrieve the top three keywords that have a score of more than 0.25
    keywords =
      predictions
      |> Enum.filter(fn keyword -> keyword.score > 0.25 end)
      |> Enum.take(3)

    send(live_view_pid, {:new_keywords, predictions})

    {:ok, %{keywords: keywords}}
  end
end