defmodule MedicalTranscriptionWeb.Components.TranscriptionTextComponent do
  @moduledoc """
  Represents a portion of transcribed text.

  Once rendered, starts tasks to find relevant codes and keywords for the text.
  """
  use MedicalTranscriptionWeb, :live_component

  import MedicalTranscriptionWeb.Components
  import MedicalTranscriptionWeb.Components.KeywordHighlighter

  alias AudioTagger.KeywordFinder
  alias MedicalTranscription.Coding
  alias MedicalTranscription.Coding.CodeVectorMatch

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

    {:ok, socket}
  end

  defp assign_initial_state(%{start_mark: start_mark, end_mark: end_mark, text: text}, socket) do
    self_pid = self()

    initial_state = %{
      start_mark: start_mark,
      end_mark: end_mark,
      text: text,
      editing: false
    }

    socket
    |> assign(initial_state)
    |> assign_async(:tags, fn -> classify_text(text) end)
    |> assign_async(:keywords, fn -> find_keywords(self_pid, text) end)
  end

  defp assign_initial_state(_assigns, socket), do: socket

  @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 !is_nil(@start_mark) && !is_nil(@end_mark) do %>
            <%= @start_mark %> - <%= @end_mark %>
          <% end %>
        </p>
        <p class="text-[28px] leading-normal text-type-black-tertiary">
          <.async_result :let={keywords} assign={@keywords}>
            <:loading><%= @text %></:loading>
            <%= if @editing do %>
              <textarea class="w-full" phx-blur="reclassify_transcription" phx-target={@myself}><%= @text %></textarea>
            <% else %>
              <div class="text-[28px]" phx-click="edit_transcription" phx-target={@myself}>
                <.highlight text={@text} keywords={keywords} />
              </div>
            <% end %>
          </.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">
        <.async_result :let={tags} assign={@tags}>
          <:loading>
            <img src={~p"/images/loading.svg"} width="16" class="animate-spin" /> Finding codes...
          </:loading>
          <:failed :let={reason}>There was an error finding codes: <%= reason %></:failed>
          <%= for %CodeVectorMatch{id: id, code: code, description: label, cosine_similarity: score, weighting: weighting} <- tags do %>
            <.tag_result
              code_vector_id={id}
              code={code}
              label={label}
              score={score}
              text={@text}
              weighting={weighting}
            />
          <% end %>
        </.async_result>
      </div>
    </div>
    """
  end

  @impl Phoenix.LiveComponent
  def handle_event("edit_transcription", _params, socket) do
    {:noreply, assign(socket, :editing, true)}
  end

  @impl Phoenix.LiveComponent
  def handle_event("reclassify_transcription", %{"value" => value} = _params, socket) do
    self_pid = self()

    socket =
      socket
      |> assign(:text, value)
      |> assign(:editing, false)
      |> assign_async(:tags, fn -> classify_text(value) end)
      |> assign_async(:keywords, fn -> find_keywords(self_pid, value) end)

    {:noreply, socket}
  end

  def classify_text(value) do
    {:ok, %{tags: Coding.process_chunk(value)}}
  end

  defp find_keywords(_live_view_pid, ""), do: {:ok, %{keywords: []}}

  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(MedicalTranscription.TokenClassificationServing, text)

    phrases = 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 = 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 = 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