/* 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,
  isMembershipFilterType,
  type MembershipFilterType,
} 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 filterType = column.columnDef.meta?.filterType;
  const membershipFilterType: MembershipFilterType | undefined =
    column.getCanFilter() && filterType && isMembershipFilterType(filterType)
      ? filterType
      : undefined;

  const handleFilterCell = (
    type: MembershipFilterType,
    operator: "in" | "not_in",
  ) => {
    const filter =
      type === "number"
        ? Filter.number({ operator, values: [rawValue] })
        : Filter.text({ operator, values: [rawValue] });
    column.setFilterValue(filter);
  };

  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>
      )}
      {membershipFilterType && (
        <>
          <ContextMenuSeparator />
          <ContextMenuItem
            onClick={() => handleFilterCell(membershipFilterType, "in")}
          >
            <FilterIcon className="mo-dropdown-icon h-3 w-3" />
            Filter by this value
          </ContextMenuItem>
          <ContextMenuItem
            onClick={() => handleFilterCell(membershipFilterType, "not_in")}
          >
            <FilterIcon className="mo-dropdown-icon h-3 w-3" />
            Remove rows with this value
          </ContextMenuItem>
        </>
      )}
    </ContextMenuContent>
  );
};