kino_table_input

defmodule KinoTableInput do
  @moduledoc """
  A table input widget for Livebook.
  """

  use Kino.JS, assets_path: "lib/assets"
  use Kino.JS.Live
  use Kino.SmartCell, name: "Table input"

  @default_table %{
    "col1" => [1, 2],
    "col2" => [1.1, 2.2],
    "col3" => ["foo", "bar"]
  }

  @default_name "table"

  @impl true
  def init(attrs, ctx) do
    table = attrs["table"] || @default_table
    columns = attrs["columns"] || Enum.map(table, fn {col, _} -> col end)
    name = attrs["name"] || @default_name
    {:ok, assign(ctx, table: table, columns: columns, name: name)}
  end

  @impl true
  def handle_connect(ctx) do
    {:ok, %{table: ctx.assigns.table, columns: ctx.assigns.columns, name: ctx.assigns.name}, ctx}
  end

  @impl true
  def handle_event("add-empty-row", %{}, ctx) do
    table =
      ctx.assigns.columns
      |> Enum.reduce(ctx.assigns.table, fn col, table ->
        Map.update!(table, col, fn rows -> rows ++ [nil] end)
      end)

    broadcast_event(ctx, "data-updated", %{
      columns: ctx.assigns.columns,
      table: table
    })

    {:noreply, assign(ctx, table: table)}
  end

  def handle_event("add-empty-column", %{}, ctx) do
    nrows =
      case Map.keys(ctx.assigns.table) do
        [key | _] -> ctx.assigns.table[key] |> length()
        [] -> 0
      end

    col_name =
      1..100
      |> Enum.find_value(fn idx ->
        col = "col#{idx}"

        if not Map.has_key?(ctx.assigns.table, col) do
          col
        end
      end)

    columns = ctx.assigns.columns ++ [col_name]
    table = Map.put(ctx.assigns.table, col_name, 0..(nrows - 1) |> Enum.map(fn _ -> nil end))

    broadcast_event(ctx, "data-updated", %{
      columns: columns,
      table: table
    })

    {:noreply, assign(ctx, table: table, columns: columns)}
  end

  def handle_event("delete-row", %{"row" => row}, ctx) do
    table =
      ctx.assigns.columns
      |> Enum.reduce(ctx.assigns.table, fn col, table ->
        Map.update!(table, col, fn rows ->
          List.delete_at(rows, row)
        end)
      end)

    broadcast_event(ctx, "data-updated", %{
      columns: ctx.assigns.columns,
      table: table
    })

    {:noreply, assign(ctx, table: table)}
  end

  def handle_event("rename-column", %{"old" => col, "new" => new_col}, ctx) do
    new_col = String.trim(new_col)

    if Map.has_key?(ctx.assigns.table, new_col) do
      broadcast_event(ctx, "data-updated", %{
        columns: ctx.assigns.columns,
        table: ctx.assigns.table
      })

      {:noreply, ctx}
    else
      columns =
        ctx.assigns.columns
        |> Enum.map(fn c -> if c == col, do: new_col, else: c end)

      table =
        case Map.get(ctx.assigns.table, col) do
          nil ->
            ctx.assigns.table

          col_vals ->
            ctx.assigns.table
            |> Map.delete(col)
            |> Map.put(new_col, col_vals)
        end

      broadcast_event(ctx, "data-updated", %{
        columns: columns,
        table: table
      })

      {:noreply, assign(ctx, table: table, columns: columns)}
    end
  end

  def handle_event("delete-column", %{"column" => col}, ctx) do
    columns = Enum.reject(ctx.assigns.columns, &(&1 == col))
    table = Map.delete(ctx.assigns.table, col)

    broadcast_event(ctx, "data-updated", %{
      columns: columns,
      table: table
    })

    {:noreply, assign(ctx, table: table, columns: columns)}
  end

  def handle_event("change-cell", %{"column" => col, "row" => row, "value" => value}, ctx) do
    table =
      case Map.get(ctx.assigns.table, col) do
        nil ->
          ctx.assigns.table

        rows ->
          rows =
            rows
            |> Enum.with_index()
            |> Enum.map(fn {x, idx} ->
              if idx == row do
                parse_value(value)
              else
                x
              end
            end)

          Map.put(ctx.assigns.table, col, rows)
      end

    broadcast_event(ctx, "data-updated", %{
      columns: ctx.assigns.columns,
      table: table
    })

    {:noreply, assign(ctx, table: table)}
  end

  def handle_event("change-name", %{"name" => name}, ctx) do
    broadcast_event(ctx, "name-updated", %{name: name})
    {:noreply, assign(ctx, name: name)}
  end

  @impl true
  def to_attrs(ctx) do
    %{
      "table" => ctx.assigns.table,
      "columns" => ctx.assigns.columns,
      "name" => ctx.assigns.name
    }
  end

  @impl true
  def to_source(attrs) do
    table =
      attrs["columns"]
      |> Enum.map(fn col ->
        {col, attrs["table"][col]}
      end)

    name = attrs["name"] || @default_name

    quote do
      unquote({String.to_atom(name), [], Elixir}) =
        unquote(table)
        |> Explorer.DataFrame.new()

      Kino.nothing()
    end
    |> Kino.SmartCell.quoted_to_string()
  end

  defp parse_value(str) do
    str = String.trim(str)

    if str == "" do
      nil
    else
      try do
        String.to_integer(str)
      rescue
        ArgumentError ->
          try do
            String.to_float(str)
          rescue
            ArgumentError -> str
          end
      end
    end
  end
end