/* Copyright 2026 Marimo. All rights reserved. */

import type { Cell } from "@tanstack/react-table";
import { useAtomValue } from "jotai";
import { CopyIcon, FilterIcon, SquareStack } from "lucide-react";
import type { RefObject } from "react";
import useEvent from "react-use-event-hook";
import { copyToClipboard } from "@/utils/copy";
import { Logger } from "@/utils/Logger";
import {
  ContextMenu,
  ContextMenuContent,
  ContextMenuItem,
  ContextMenuPortal,
  ContextMenuSeparator,
  ContextMenuTrigger,
} from "../ui/context-menu";
import { DATA_CELL_ID } from "./cell-utils";
import { Filter } from "./filters";
import { selectedCellsAtom } from "./range-focus/atoms";
import { getClipboardContent, getRawValue } from "./utils";

export const DataTableContextMenu = <TData,>({
  contextMenuRef,
  tableBody,
  tableRef,
  copyAllCells,
}: {
  contextMenuRef: RefObject<Cell<TData, unknown> | null>;
  tableBody: React.ReactNode;
  tableRef: RefObject<HTMLTableSectionElement | null>;
  copyAllCells: () => void;
}) => {
  const handleContextMenuChange = useEvent((open: boolean) => {
    const cell = contextMenuRef.current;
    if (!cell) {
      return;
    }

    // Add a background color to the cell when the context menu is open
    const cellElement = tableRef.current?.querySelector(
      `[${DATA_CELL_ID}="${cell.id}"]`,
    );
    if (!cellElement) {
      Logger.error("Context menu cell not found in table");
      return;
    }

    if (open) {
      cellElement.classList.add("bg-(--green-4)");
    } else {
      cellElement.classList.remove("bg-(--green-4)");
    }
  });

  return (
    <ContextMenu onOpenChange={handleContextMenuChange}>
      <ContextMenuTrigger asChild={true}>{tableBody}</ContextMenuTrigger>
      <ContextMenuPortal>
        <CellContextMenu
          cellRef={contextMenuRef}
          copySelectedCells={copyAllCells}
        />
      </ContextMenuPortal>
    </ContextMenu>
  );
};

export const CellContextMenu = <TData,>({
  cellRef,
  copySelectedCells,
}: {
  cellRef: RefObject<Cell<TData, unknown> | null>;
  copySelectedCells: () => void;
}) => {
  const selectedCells = useAtomValue(selectedCellsAtom);
  const multipleSelectedCells = selectedCells.size > 1;

  const cell = cellRef.current;
  if (!cell) {
    Logger.error("No cell found in context menu");
    return;
  }

  const table = cell.getContext().table;
  const displayedValue = cell.getValue();
  const rawValue =
    getRawValue(table, cell.row.index, cell.column.id) ?? displayedValue;

  const handleCopyCell = () => {
    try {
      const { text, html } = getClipboardContent(rawValue, displayedValue);
      copyToClipboard(text, html);
    } catch (error) {
      Logger.error("Failed to copy context menu cell", error);
    }
  };

  const column = cell.column;
  const canFilter = column.getCanFilter() && column.columnDef.meta?.filterType;

  const handleFilterCell = (operator: "in" | "not_in") => {
    column.setFilterValue(
      Filter.select({
        options: [rawValue],
        operator,
      }),
    );
  };

  return (
    <ContextMenuContent>
      <ContextMenuItem onClick={handleCopyCell}>
        <CopyIcon className="mo-dropdown-icon h-3 w-3" />
        Copy cell
      </ContextMenuItem>
      {multipleSelectedCells && (
        <ContextMenuItem onClick={copySelectedCells}>
          <SquareStack className="mo-dropdown-icon h-3 w-3" />
          Copy selected cells
        </ContextMenuItem>
      )}
      {canFilter && (
        <>
          <ContextMenuSeparator />
          <ContextMenuItem onClick={() => handleFilterCell("in")}>
            <FilterIcon className="mo-dropdown-icon h-3 w-3" />
            Filter by this value
          </ContextMenuItem>
          <ContextMenuItem onClick={() => handleFilterCell("not_in")}>
            <FilterIcon className="mo-dropdown-icon h-3 w-3" />
            Remove rows with this value
          </ContextMenuItem>
        </>
      )}
    </ContextMenuContent>
  );
};