arcana

defmodule ArcanaWeb.AskLive do
  @moduledoc """
  LiveView for asking questions about documents in Arcana.
  """
  use Phoenix.LiveView

  import Ecto.Query
  import ArcanaWeb.DashboardComponents

  alias Arcana.Graph.Entity

  @impl true
  def mount(_params, session, socket) do
    repo = get_repo_from_session(session)

    {:ok,
     socket
     |> assign(repo: repo)
     |> assign(
       ask_mode: :agentic,
       ask_question: "",
       ask_running: false,
       ask_context: nil,
       ask_error: nil,
       stats: nil,
       collections: [],
       graph_search: true,
       selected_collections: [],
       graph_context_expanded: true,
       llm_select: false,
       pipeline_step: nil
     )}
  end

  @impl true
  def handle_params(_params, _uri, socket) do
    {:noreply, load_data(socket)}
  end

  defp load_data(socket) do
    repo = socket.assigns.repo

    socket
    |> assign(stats: load_stats(repo))
    |> assign(collections: load_collections_with_graph_status(repo))
  end

  defp load_collections_with_graph_status(repo) do
    collections = load_collections(repo)

    # Get entity counts per collection
    entity_counts =
      repo.all(
        from(e in Entity,
          group_by: e.collection_id,
          select: {e.collection_id, count(e.id)}
        )
      )
      |> Map.new()

    Enum.map(collections, fn c ->
      entity_count = Map.get(entity_counts, c.id, 0)
      Map.put(c, :graph_enabled, entity_count > 0)
    end)
  end

  defp selected_graph_enabled?(collections, selected) do
    case selected do
      [] -> false
      names -> Enum.any?(collections, &(&1.name in names and &1[:graph_enabled]))
    end
  end

  @impl true
  def handle_event("ask_submit", params, socket) do
    question = params["question"] || ""

    case {Application.get_env(:arcana, :llm), question} do
      {nil, _} ->
        {:noreply,
         assign(socket,
           ask_error: "No LLM configured. Set :arcana, :llm in your config.",
           ask_running: false
         )}

      {_, ""} ->
        {:noreply, assign(socket, ask_error: "Please enter a question")}

      {llm, question} ->
        socket = assign(socket, ask_running: true, ask_error: nil, ask_question: question)
        start_ask_task(socket, params, llm, question)
        {:noreply, socket}
    end
  end

  def handle_event("ask_clear", _params, socket) do
    {:noreply, assign(socket, ask_context: nil, ask_error: nil, ask_question: "")}
  end

  def handle_event("ask_switch_mode", %{"mode" => mode}, socket) do
    mode = String.to_existing_atom(mode)
    # Reset llm_select when switching to simple mode
    llm_select = if mode == :simple, do: false, else: socket.assigns.llm_select
    {:noreply, assign(socket, ask_mode: mode, llm_select: llm_select)}
  end

  def handle_event("form_changed", params, socket) do
    selected = params["collections"] || []
    {:noreply, assign(socket, selected_collections: selected)}
  end

  def handle_event("toggle_graph_context", _params, socket) do
    {:noreply, assign(socket, graph_context_expanded: !socket.assigns.graph_context_expanded)}
  end

  def handle_event("toggle_llm_select", _params, socket) do
    {:noreply, assign(socket, llm_select: !socket.assigns.llm_select)}
  end

  defp start_ask_task(socket, params, llm, question) do
    repo = socket.assigns.repo
    mode = params["mode"] || "simple"
    selected_collections = params["collections"] || []
    parent = self()

    Arcana.TaskSupervisor.start_child(fn ->
      handler_id = "pipeline-progress-#{inspect(parent)}"

      graph_enabled = params["graph_search"] == "true"

      steps = [:expand, :decompose, :select, :search, :self_correct, :rerank, :answer, :ground]

      events =
        Enum.map(steps, &[:arcana, :agent, &1, :start]) ++
          [[:arcana, :graph, :search, :start]]

      :telemetry.attach_many(
        handler_id,
        events,
        fn
          [:arcana, :graph, :search, :start], _measurements, _metadata, _config ->
            send(parent, {:pipeline_progress, "Searching with graph connections..."})

          [:arcana, :agent, step, :start], _measurements, _metadata, _config ->
            label =
              if step == :search and graph_enabled,
                do: "Searching with graph connections...",
                else: pipeline_step_label(step)

            if label, do: send(parent, {:pipeline_progress, label})
        end,
        nil
      )

      result =
        run_ask(
          mode,
          question,
          repo,
          llm,
          socket.assigns.collections,
          params,
          selected_collections
        )

      :telemetry.detach(handler_id)
      send(parent, {:ask_complete, result})
    end)
  end

  defp run_ask("simple", question, repo, llm, _all_collections, params, selected_collections) do
    run_simple_ask(question, repo, llm, selected_collections, params)
  end

  defp run_ask(_mode, question, repo, llm, all_collections, params, selected_collections) do
    run_agentic_ask(
      question,
      repo,
      llm,
      all_collections,
      collections: selected_collections,
      use_llm_select: params["llm_select"] == "true",
      use_expand: params["use_expand"] == "true",
      use_decompose: params["use_decompose"] == "true",
      use_rerank: params["use_rerank"] == "true",
      use_ground: params["use_ground"] == "true",
      self_correct: params["self_correct"] == "true",
      hallucinate_demo: params["answer_mode"] == "hallucinate",
      graph: params["graph_search"] == "true"
    )
  end

  @impl true
  def handle_info({:pipeline_progress, step}, socket) do
    {:noreply, assign(socket, pipeline_step: step)}
  end

  def handle_info({:ask_complete, result}, socket) do
    socket =
      case result do
        {:ok, ctx} ->
          assign(socket, ask_running: false, ask_context: ctx, ask_error: nil, pipeline_step: nil)

        {:error, reason} ->
          assign(socket, ask_running: false, ask_error: inspect(reason), pipeline_step: nil)
      end

    {:noreply, socket}
  end

  defp run_simple_ask(question, repo, llm, selected_collections, params) do
    graph = params["graph_search"] == "true"
    opts = [repo: repo, llm: llm, graph: graph]

    opts =
      case selected_collections do
        [] -> opts
        [single] -> Keyword.put(opts, :collection, single)
        multiple -> Keyword.put(opts, :collections, multiple)
      end

    case Arcana.ask(question, opts) do
      {:ok, answer, results} ->
        # Build a context-like struct for consistent UI display
        {:ok,
         %{
           question: question,
           answer: answer,
           results: results,
           expanded_query: nil,
           sub_questions: nil,
           selected_collections:
             if(selected_collections == [], do: nil, else: selected_collections)
         }}

      {:error, reason} ->
        {:error, reason}
    end
  rescue
    e -> {:error, Exception.message(e)}
  end

  defp run_agentic_ask(question, repo, llm, all_collections, opts) do
    alias Arcana.Agent

    all_collection_names = Enum.map(all_collections, & &1.name)
    search_opts = build_search_opts(opts, all_collection_names)

    Agent.new(question, repo: repo, llm: llm)
    |> maybe_select(opts, all_collection_names)
    |> maybe_expand(opts)
    |> maybe_decompose(opts)
    |> Agent.search(search_opts)
    |> maybe_rerank(opts)
    |> maybe_answer_with_hallucinations(opts)
    |> maybe_ground(opts)
    |> format_agentic_result(question)
  rescue
    e -> {:error, Exception.message(e)}
  end

  defp maybe_select(ctx, opts, all_collection_names) do
    if Keyword.get(opts, :use_llm_select, false) and length(all_collection_names) > 1 do
      Arcana.Agent.select(ctx, collections: all_collection_names)
    else
      ctx
    end
  end

  defp maybe_expand(ctx, opts) do
    if Keyword.get(opts, :use_expand, false), do: Arcana.Agent.expand(ctx), else: ctx
  end

  defp maybe_decompose(ctx, opts) do
    if Keyword.get(opts, :use_decompose, false), do: Arcana.Agent.decompose(ctx), else: ctx
  end

  defp maybe_rerank(ctx, opts) do
    if Keyword.get(opts, :use_rerank, false), do: Arcana.Agent.rerank(ctx), else: ctx
  end

  defp maybe_answer_with_hallucinations(ctx, opts) do
    if Keyword.get(opts, :hallucinate_demo, false) do
      Arcana.Agent.answer(ctx, prompt: &hallucination_demo_prompt/2)
    else
      Arcana.Agent.answer(ctx)
    end
  end

  defp hallucination_demo_prompt(question, chunks) do
    reference_material =
      chunks
      |> Enum.with_index(1)
      |> Enum.map_join("\n\n", fn {chunk, i} -> "[#{i}] #{chunk.text}" end)

    """
    Context:
    #{reference_material}

    Question: "#{question}"

    Answer the question using the context above, but deliberately slip in 1-2 plausible-sounding statements that are NOT supported by the context. These fabricated facts should blend naturally into the answer. Do not flag or mark which statements are made up.
    """
  end

  defp maybe_ground(ctx, opts) do
    if Keyword.get(opts, :use_ground, false), do: Arcana.Agent.ground(ctx), else: ctx
  end

  defp build_search_opts(opts, all_collection_names) do
    base = [
      self_correct: Keyword.get(opts, :self_correct, false),
      graph: Keyword.get(opts, :graph, false)
    ]

    use_llm_select = Keyword.get(opts, :use_llm_select, false)

    if use_llm_select and length(all_collection_names) > 1 do
      base
    else
      add_collection_opts(base, Keyword.get(opts, :collections, []))
    end
  end

  defp add_collection_opts(opts, []), do: opts
  defp add_collection_opts(opts, [single]), do: Keyword.put(opts, :collection, single)
  defp add_collection_opts(opts, multiple), do: Keyword.put(opts, :collections, multiple)

  defp format_agentic_result(%{error: error}, _question) when not is_nil(error) do
    {:error, error}
  end

  defp format_agentic_result(ctx, question) do
    # Flatten chunks from nested results to match simple mode format
    all_chunks =
      (ctx.results || [])
      |> Enum.flat_map(fn
        %{chunks: chunks} -> chunks
        chunk -> [chunk]
      end)
      |> Enum.uniq_by(& &1.id)

    {:ok,
     %{
       question: question,
       answer: ctx.answer,
       results: all_chunks,
       expanded_query: ctx.expanded_query,
       sub_questions: ctx.sub_questions,
       selected_collections: ctx.collections,
       grounding: ctx.grounding
     }}
  end

  @impl true
  def render(assigns) do
    ~H"""
    <.dashboard_layout stats={@stats} current_tab={:ask}>
      <div class="arcana-ask">
        <h2>Ask</h2>
        <p class="arcana-tab-description">
          Ask questions about your documents. Choose Simple for basic RAG or Agentic for advanced pipeline features.
        </p>

        <div class="arcana-ask-mode-nav">
          <button
            class={"arcana-mode-btn #{if @ask_mode == :agentic, do: "active", else: ""}"}
            phx-click="ask_switch_mode"
            phx-value-mode="agentic"
          >
            Agentic
          </button>
          <button
            class={"arcana-mode-btn #{if @ask_mode == :simple, do: "active", else: ""}"}
            phx-click="ask_switch_mode"
            phx-value-mode="simple"
          >
            Simple
          </button>
        </div>

        <p class="arcana-mode-description">
          <%= if @ask_mode == :simple do %>
            Basic RAG: search for relevant chunks and generate an answer.
          <% else %>
            Advanced RAG: query expansion, decomposition, self-correction, and reranking.
          <% end %>
        </p>

        <%= if @ask_error do %>
          <div class="arcana-eval-message error">
            <%= @ask_error %>
          </div>
        <% end %>

        <form id="ask-form" phx-submit="ask_submit" phx-change="form_changed" class="arcana-ask-form">
          <input type="hidden" name="mode" value={@ask_mode} />

          <div class="arcana-ask-input">
            <textarea
              name="question"
              placeholder="Ask a question about your documents..."
              rows="3"
              disabled={@ask_running}
            ><%= @ask_question %></textarea>

            <%= if @ask_mode == :simple and selected_graph_enabled?(@collections, @selected_collections) do %>
              <label class="arcana-deep-search-toggle">
                <input
                  type="checkbox"
                  name="graph_search"
                  value="true"
                  checked={@graph_search}
                  disabled={@ask_running}
                />
                <span>Graph-Assisted</span>
                <small>Find results through entity relationships</small>
              </label>
            <% end %>
          </div>

          <div class="arcana-ask-collections">
            <label>Collections</label>
            <%= if @ask_mode == :agentic and length(@collections) > 1 do %>
              <div class="arcana-llm-select-toggle">
                <label class="arcana-checkbox-label">
                  <input
                    type="checkbox"
                    name="llm_select"
                    value="true"
                    checked={@llm_select}
                    disabled={@ask_running}
                    phx-click="toggle_llm_select"
                  />
                  <span>Let LLM select automatically</span>
                  <small>LLM will choose the most relevant collection(s) based on your question</small>
                </label>
              </div>
            <% end %>
            <%= if not @llm_select do %>
              <div class="arcana-collection-checkboxes">
                <%= for coll <- @collections do %>
                  <label class="arcana-collection-check">
                    <input
                      type="checkbox"
                      name="collections[]"
                      value={coll.name}
                      checked={coll.name in @selected_collections}
                      disabled={@ask_running}
                    />
                    <span><%= coll.name %></span>
                  </label>
                <% end %>
              </div>
              <small class="arcana-collection-hint">Select none for all collections</small>
            <% end %>
          </div>

          <%= if @ask_mode == :agentic do %>
            <div class="arcana-ask-options">
              <h4>
                Pipeline
                <span style="font-size: 0.75em; font-weight: normal; opacity: 0.6;">
                  <a href="#" onclick="this.closest('.arcana-ask-options').querySelectorAll('input[type=checkbox]').forEach(c => c.checked = true); return false">all</a>
                  /
                  <a href="#" onclick="this.closest('.arcana-ask-options').querySelectorAll('input[type=checkbox]').forEach(c => c.checked = false); return false">none</a>
                </span>
              </h4>
              <ol class="arcana-pipeline">
                <li>
                  <label class="arcana-pipeline-step">
                    <input type="checkbox" name="use_expand" value="true" disabled={@ask_running} />
                    <span class="arcana-step-label">Query Expansion</span>
                    <small>Generate related queries</small>
                  </label>
                </li>
                <li>
                  <label class="arcana-pipeline-step">
                    <input type="checkbox" name="use_decompose" value="true" disabled={@ask_running} />
                    <span class="arcana-step-label">Decomposition</span>
                    <small>Break into sub-questions</small>
                  </label>
                </li>
                <li>
                  <%= if selected_graph_enabled?(@collections, @selected_collections) do %>
                    <div class="arcana-pipeline-fork">
                      <label class="arcana-pipeline-step">
                        <input type="radio" name="graph_search" value="false" disabled={@ask_running} />
                        <span class="arcana-step-label">Search</span>
                        <small>Retrieve relevant chunks</small>
                      </label>
                      <span class="arcana-fork-or">or</span>
                      <label class="arcana-pipeline-step">
                        <input type="radio" name="graph_search" value="true" checked disabled={@ask_running} />
                        <span class="arcana-step-label">Graph-Assisted Search</span>
                        <small>Find results through entity relationships</small>
                      </label>
                    </div>
                  <% else %>
                    <div class="arcana-pipeline-step fixed">
                      <span class="arcana-step-label">Search</span>
                      <small>Retrieve relevant chunks</small>
                    </div>
                  <% end %>
                </li>
                <li>
                  <label class="arcana-pipeline-step">
                    <input type="checkbox" name="self_correct" value="true" disabled={@ask_running} />
                    <span class="arcana-step-label">Self-Correction</span>
                    <small>Refine search if results are poor</small>
                  </label>
                </li>
                <li>
                  <label class="arcana-pipeline-step">
                    <input type="checkbox" name="use_rerank" value="true" disabled={@ask_running} />
                    <span class="arcana-step-label">Reranking</span>
                    <small>LLM-based result reranking</small>
                  </label>
                </li>
                <li>
                  <div class="arcana-pipeline-fork">
                    <label class="arcana-pipeline-step">
                      <input type="radio" name="answer_mode" value="normal" checked disabled={@ask_running} />
                      <span class="arcana-step-label">Answer</span>
                      <small>Generate faithful answer</small>
                    </label>
                    <span class="arcana-fork-or">or</span>
                    <label class="arcana-pipeline-step">
                      <input type="radio" name="answer_mode" value="hallucinate" disabled={@ask_running} />
                      <span class="arcana-step-label">Answer with Hallucination</span>
                      <small>Mix in fake facts to showcase grounding</small>
                    </label>
                  </div>
                </li>
                <li>
                  <label class="arcana-pipeline-step">
                    <input type="checkbox" name="use_ground" value="true" disabled={@ask_running} />
                    <span class="arcana-step-label">Grounding</span>
                    <small>Detect hallucinated vs faithful spans</small>
                  </label>
                </li>
              </ol>
            </div>
          <% end %>

          <div class="arcana-ask-actions">
            <button type="submit" disabled={@ask_running}>
              <%= if @ask_running, do: "Asking...", else: "Ask" %>
            </button>
            <%= if @ask_context do %>
              <button type="button" phx-click="ask_clear" disabled={@ask_running}>
                Clear
              </button>
            <% end %>
          </div>
        </form>

        <%= if @ask_running do %>
          <div class="arcana-ask-loading">
            <div class="arcana-spinner"></div>
            <span><%= if @ask_mode == :simple, do: "Generating answer...", else: @pipeline_step || "Running pipeline..." %></span>
          </div>
        <% end %>

        <%= if @ask_context do %>
          <div class="arcana-ask-results">
            <div class="arcana-ask-answer">
              <h3>Answer</h3>
              <div class="arcana-answer-content">
                <%= if @ask_context.answer do %>
                  <%= if Map.get(@ask_context, :grounding) do %>
                    <%= render_highlighted_answer(@ask_context.answer, @ask_context.grounding) %>
                  <% else %>
                    <%= @ask_context.answer %>
                  <% end %>
                <% else %>
                  <span style="color: #9ca3af; font-style: italic;">No answer generated</span>
                <% end %>
              </div>
            </div>

            <%= if Map.get(@ask_context, :grounding) do %>
              <div class="arcana-ask-section arcana-grounding-results">
                <h4>
                  Grounding
                  <span class={"arcana-grounding-score #{if @ask_context.grounding.score >= 0.8, do: "good", else: if(@ask_context.grounding.score >= 0.5, do: "warn", else: "bad")}"}>
                    <%= Float.round(@ask_context.grounding.score * 100, 1) %>% faithful
                  </span>
                </h4>

                <% chunks = Map.get(@ask_context, :results, []) %>

                <%= if length(@ask_context.grounding.hallucinated_spans) > 0 do %>
                  <div class="arcana-grounding-spans">
                    <h5>Hallucinated Spans</h5>
                    <%= for span <- @ask_context.grounding.hallucinated_spans do %>
                      <div class="arcana-span hallucinated">
                        <span class="arcana-span-text"><%= span.text %></span>
                        <span class="arcana-span-score"><%= Float.round(span.score, 2) %></span>
                        <.source_previews sources={Map.get(span, :sources, [])} chunks={chunks} />
                      </div>
                    <% end %>
                  </div>
                <% end %>

                <%= if length(@ask_context.grounding.faithful_spans) > 0 do %>
                  <div class="arcana-grounding-spans">
                    <h5>Faithful Spans (<%= length(@ask_context.grounding.faithful_spans) %>)</h5>
                    <%= for span <- @ask_context.grounding.faithful_spans do %>
                      <div class="arcana-span faithful">
                        <span class="arcana-span-text"><%= span.text %></span>
                        <.source_previews sources={Map.get(span, :sources, [])} chunks={chunks} />
                      </div>
                    <% end %>
                  </div>
                <% end %>
              </div>
            <% end %>

            <%= if @ask_context.expanded_query do %>
              <div class="arcana-ask-section">
                <h4>Expanded Query</h4>
                <p class="arcana-expanded-query"><%= @ask_context.expanded_query %></p>
              </div>
            <% end %>

            <%= if @ask_context.sub_questions && length(@ask_context.sub_questions) > 0 do %>
              <div class="arcana-ask-section">
                <h4>Sub-Questions</h4>
                <ul class="arcana-query-list">
                  <%= for sq <- @ask_context.sub_questions do %>
                    <li><%= sq %></li>
                  <% end %>
                </ul>
              </div>
            <% end %>

            <%= if @ask_context.selected_collections && length(@ask_context.selected_collections) > 0 do %>
              <div class="arcana-ask-section">
                <h4>Selected Collections</h4>
                <div class="arcana-collection-badges">
                  <%= for coll <- @ask_context.selected_collections do %>
                    <span class="arcana-collection-badge"><%= coll %></span>
                  <% end %>
                </div>
              </div>
            <% end %>

            <%= if Map.get(@ask_context, :graph_enhanced) do %>
              <.graph_context_section
                matched_entities={Map.get(@ask_context, :matched_entities, [])}
                matched_relationships={Map.get(@ask_context, :matched_relationships, [])}
                expanded={@graph_context_expanded}
              />
            <% end %>

            <%= if @ask_context.results && length(@ask_context.results) > 0 do %>
              <% all_chunks = @ask_context.results %>
              <div class="arcana-ask-section">
                <h4>Retrieved Chunks (<%= length(all_chunks) %>)</h4>
                <div class="arcana-search-results">
                  <%= for chunk <- all_chunks do %>
                    <div class="arcana-search-result">
                      <div class="arcana-result-header">
                        <div class="arcana-result-score">
                          <span class="score-value"><%= Float.round(chunk.score, 4) %></span>
                          <%= if Map.get(chunk, :graph_sources) && length(chunk.graph_sources) > 0 do %>
                            <span class="arcana-graph-attribution">
                              via: <%= Enum.join(chunk.graph_sources, ", ") %>
                            </span>
                          <% end %>
                        </div>
                        <div class="arcana-result-meta">
                          <code><%= chunk.document_id %></code>
                          <span class="arcana-chunk-badge">Chunk <%= chunk.chunk_index %></span>
                        </div>
                      </div>
                      <div class="arcana-result-text">
                        <%= String.slice(chunk.text, 0, 300) %><%= if String.length(chunk.text) > 300, do: "...", else: "" %>
                      </div>
                    </div>
                  <% end %>
                </div>
              </div>
            <% end %>
          </div>
        <% end %>
      </div>
    </.dashboard_layout>
    """
  end

  defp graph_context_section(assigns) do
    ~H"""
    <div class="arcana-graph-context">
      <div class="arcana-graph-context-header">
        <h4>Graph Context</h4>
        <button type="button" phx-click="toggle_graph_context" class="arcana-toggle-btn">
          <%= if @expanded, do: "▼", else: "▶" %>
        </button>
      </div>

      <%= if @expanded do %>
        <div class="arcana-graph-context-content">
          <%= if length(@matched_entities) == 0 and length(@matched_relationships) == 0 do %>
            <p class="arcana-no-matches">No entity matches — used vector search only</p>
          <% else %>
            <%= if length(@matched_entities) > 0 do %>
              <div class="arcana-matched-entities">
                <h5>Matched Entities</h5>
                <ul>
                  <%= for entity <- @matched_entities do %>
                    <li>
                      <span class="arcana-entity-name"><%= entity.name %></span>
                      <span class="arcana-entity-type"><%= entity.type %></span>
                      <%= if Map.get(entity, :id) do %>
                        <a href={"/arcana/graph?entity=#{entity.id}"} class="arcana-view-in-graph">
                          View in Graph
                        </a>
                      <% end %>
                    </li>
                  <% end %>
                </ul>
              </div>
            <% end %>

            <%= if length(@matched_relationships) > 0 do %>
              <div class="arcana-matched-relationships">
                <h5>Key Relationships</h5>
                <ul>
                  <%= for rel <- @matched_relationships do %>
                    <li>
                      <span class="arcana-rel-source"><%= rel.source %></span>
                      <span class="arcana-rel-type">—<%= rel.type %>→</span>
                      <span class="arcana-rel-target"><%= rel.target %></span>
                    </li>
                  <% end %>
                </ul>
              </div>
            <% end %>
          <% end %>
        </div>
      <% end %>
    </div>
    """
  end

  defp source_previews(assigns) do
    ~H"""
    <%= if @sources != [] do %>
      <div class="arcana-span-sources">
        <%= for source <- @sources do %>
          <% chunk = find_chunk(@chunks, source.chunk_id) %>
          <details class="arcana-source-detail">
            <summary class="arcana-source-badge clickable">
              <span class="arcana-source-label">
                <%= chunk_label(chunk) %>
              </span>
              <span class="arcana-source-overlap"><%= Float.round(source.score * 100) %>% overlap</span>
            </summary>
            <%= if chunk do %>
              <div class="arcana-source-preview">
                <%= chunk.text %>
              </div>
            <% end %>
          </details>
        <% end %>
      </div>
    <% end %>
    """
  end

  defp find_chunk(chunks, chunk_id) do
    Enum.find(chunks, fn chunk ->
      id = if is_struct(chunk), do: chunk.id, else: Map.get(chunk, :id)
      to_string(id) == to_string(chunk_id)
    end)
  end

  defp chunk_label(nil), do: "unknown chunk"

  defp chunk_label(chunk) do
    index = Map.get(chunk, :chunk_index, nil)
    if index, do: "Chunk #{index}", else: "Source"
  end

  defp render_highlighted_answer(answer, grounding) do
    hallucinated = Map.get(grounding, :hallucinated_spans, [])
    faithful = Map.get(grounding, :faithful_spans, [])

    spans =
      (Enum.map(hallucinated, &Map.put(&1, :type, :hallucinated)) ++
         Enum.map(faithful, &Map.put(&1, :type, :faithful)))
      |> Enum.sort_by(& &1.start)

    build_highlighted_parts(answer, spans, 0, [])
    |> Enum.reverse()
    |> Phoenix.HTML.raw()
  end

  defp build_highlighted_parts(answer, [], pos, acc) do
    rest = binary_slice(answer, pos, byte_size(answer) - pos)
    if rest == "", do: acc, else: [html_escape(rest) | acc]
  end

  defp build_highlighted_parts(answer, [span | rest], pos, acc) do
    span_start = span.start
    span_end = Map.get(span, :end, span_start)

    # Text before this span
    acc =
      if span_start > pos do
        gap = binary_slice(answer, pos, span_start - pos)
        [html_escape(gap) | acc]
      else
        acc
      end

    # The span itself
    span_text = binary_slice(answer, span_start, span_end - span_start)

    class =
      if span.type == :hallucinated, do: "arcana-hl-hallucinated", else: "arcana-hl-faithful"

    acc = [
      ~s(<mark class="#{class}" title="#{span.type}">#{html_escape(span_text)}</mark>) | acc
    ]

    build_highlighted_parts(answer, rest, span_end, acc)
  end

  defp html_escape(text) do
    text
    |> Phoenix.HTML.html_escape()
    |> Phoenix.HTML.safe_to_string()
  end

  defp pipeline_step_label(:expand), do: "Expanding query..."
  defp pipeline_step_label(:decompose), do: "Decomposing question..."
  defp pipeline_step_label(:select), do: "Selecting collections..."
  defp pipeline_step_label(:search), do: "Searching..."
  defp pipeline_step_label(:self_correct), do: "Refining search..."
  defp pipeline_step_label(:rerank), do: "Reranking results..."
  defp pipeline_step_label(:answer), do: "Generating answer..."
  defp pipeline_step_label(:ground), do: "Checking for hallucinations..."
  defp pipeline_step_label(_), do: nil
end