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

import type { Table } from "@tanstack/react-table";
import { SELECT_COLUMN_ID } from "../types";
import { getClipboardContent, getRawValue } from "../utils";
import type { SelectedCell } from "./atoms";

export interface CellValuesResult {
  text: string;
  html: string | undefined;
}

/**
 * Get the values of the selected cells, preferring raw (unformatted) values
 * for plain text. If any cell contains HTML (e.g. a hyperlink), also builds
 * an HTML table so rich content can be preserved on paste.
 */
export function getCellValues<TData>(
  table: Table<TData>,
  selectedCellIds: Set<string>,
): CellValuesResult {
  const rows = new Map<string, { text: string; html?: string }[]>();
  let hasHtml = false;

  for (const cellId of selectedCellIds) {
    if (cellId.includes(SELECT_COLUMN_ID)) {
      // Ignore select checkbox in tables
      continue;
    }

    const { rowId, columnId } = getRowAndColumnId(cellId);
    const row = table.getRow(rowId);
    if (!row) {
      continue;
    }

    const rawValue = getRawValue(table, row.index, columnId);
    const { text, html } = getClipboardContent(
      rawValue,
      row.getValue(columnId),
    );

    if (html) {
      hasHtml = true;
    }
    const cells = rows.get(rowId) ?? [];
    cells.push({ text, html });
    rows.set(rowId, cells);
  }

  const rowValues = [...rows.values()];
  const tabSeparatedText = rowValues
    .map((cells) => cells.map((c) => c.text).join("\t"))
    .join("\n");

  let htmlTable: string | undefined;
  if (hasHtml) {
    const htmlTableRows = rowValues
      .map(
        (cells) =>
          `<tr>${cells.map((c) => `<td>${c.html ?? escapeHtml(c.text)}</td>`).join("")}</tr>`,
      )
      .join("");
    htmlTable = `<table>${htmlTableRows}</table>`;
  }

  return {
    text: tabSeparatedText,
    html: htmlTable,
  };
}

function escapeHtml(str: string): string {
  const div = document.createElement("div");
  div.textContent = str;
  return div.innerHTML;
}

/**
 * Count selected cells excluding the select checkbox column.
 */
export function countDataCellsInSelection(
  selectedCellIds: Set<string>,
): number {
  let count = 0;
  for (const cellId of selectedCellIds) {
    if (!cellId.includes(SELECT_COLUMN_ID)) {
      count += 1;
    }
  }
  return count;
}

/**
 * Extract numeric values from the selected cells. Only finite numbers and
 * non-empty numeric strings (e.g. "42", "3.14", "0") are included. Skips select
 * checkbox column, missing cells, and all other types (boolean, null, etc.).
 */
export function getNumericValuesFromSelectedCells<TData>(
  table: Table<TData>,
  selectedCellIds: Set<string>,
): number[] {
  const numericValues: number[] = [];
  for (const cellId of selectedCellIds) {
    if (cellId.includes(SELECT_COLUMN_ID)) {
      continue;
    }
    const { rowId, columnId } = getRowAndColumnId(cellId);
    const row = table.getRow(rowId);
    if (!row) {
      continue;
    }

    const value =
      getRawValue(table, row.index, columnId) ?? row.getValue(columnId);

    // Only accept numbers and strings
    // Skip booleans, null, etc.
    let num: number;
    if (typeof value === "number") {
      num = value;
    } else if (typeof value === "string") {
      if (value.trim() === "") {
        continue;
      }
      num = Number(value);
    } else {
      continue;
    }

    // Skip NaN and Infinity
    if (Number.isFinite(num)) {
      numericValues.push(num);
    }
  }
  return numericValues;
}

/**
 * Get the cell ids between two cells.
 */
export function getCellsBetween<TData>(
  table: Table<TData>,
  startCell: SelectedCell,
  endCell: SelectedCell,
): string[] {
  const startRow = table.getRow(startCell.rowId);
  const endRow = table.getRow(endCell.rowId);

  if (!startRow || !endRow) {
    return [];
  }

  const startRowIdx = startRow.index;
  const endRowIdx = endRow.index;
  const startColumnIdx = table.getColumn(startCell.columnId)?.getIndex();
  const endColumnIdx = table.getColumn(endCell.columnId)?.getIndex();

  if (startColumnIdx === undefined || endColumnIdx === undefined) {
    return [];
  }

  const minRow = Math.min(startRowIdx, endRowIdx);
  const maxRow = Math.max(startRowIdx, endRowIdx);
  const minCol = Math.min(startColumnIdx, endColumnIdx);
  const maxCol = Math.max(startColumnIdx, endColumnIdx);

  // Pre-allocate array with known size
  const result: string[] = [];
  const totalCells = (maxRow - minRow + 1) * (maxCol - minCol + 1);
  result.length = totalCells;
  let resultIndex = 0;

  const columnIds = table.getAllColumns().map((col) => col.id);
  const rows = table.getRowModel().rows;

  for (let i = minRow; i <= maxRow; i++) {
    const row = rows[i];
    const rowId = row.id;

    for (let j = minCol; j <= maxCol; j++) {
      const columnId = columnIds[j];
      result[resultIndex++] = getCellId(rowId, columnId);
    }
  }

  // Trim any unused slots
  result.length = resultIndex;
  return result;
}

/**
 * By default, the cell id is the row id and the column id separated by an underscore.
 * https://tanstack.com/table/latest/docs/guide/cells#cell-ids
 */
function getCellId(rowId: string, columnId: string) {
  return `${rowId}_${columnId}`;
}

function getRowAndColumnId(cellId: string) {
  const sepIdx = cellId.indexOf("_");
  const rowId = cellId.slice(0, sepIdx);
  const columnId = cellId.slice(sepIdx + 1);
  return { rowId, columnId };
}