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

import type { Table } from "@tanstack/react-table";
import { render, screen } from "@testing-library/react";
import { useEffect } from "react";
import { I18nProvider } from "react-aria";
import { describe, expect, it } from "vitest";
import { TooltipProvider } from "@/components/ui/tooltip";
import { SELECT_COLUMN_ID } from "../../types";
import { useCellSelectionReducerActions } from "../atoms";
import { CellSelectionStats } from "../cell-selection-stats";
import { CellSelectionProvider } from "../provider";
import { createMockCell, createMockRow, createMockTable } from "./test-utils";

const TestHarness = ({
  table,
  selectedCellIds,
}: {
  table: Table<unknown>;
  selectedCellIds: Set<string>;
}) => {
  const actions = useCellSelectionReducerActions();
  useEffect(() => {
    actions.setSelectedCells(selectedCellIds);
  }, [actions, selectedCellIds]);
  return (
    <TooltipProvider>
      <I18nProvider locale="en-US">
        <CellSelectionStats table={table} />
      </I18nProvider>
    </TooltipProvider>
  );
};

describe("CellSelectionStats", () => {
  it("should show hint when fewer than 2 cells are selected", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 10),
      createMockCell("0_1", 20),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("No selection")).toBeInTheDocument();
    expect(screen.queryByText(/Count:/)).not.toBeInTheDocument();
  });

  it("should display Count stat when 2 or more cells are selected", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 10),
      createMockCell("0_1", 20),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0", "0_1"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Count: 2")).toBeInTheDocument();
  });

  it("should display Sum when selection contains numeric values", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 10),
      createMockCell("0_1", 20),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0", "0_1"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Sum: 30")).toBeInTheDocument();
  });

  it("should display Average when selection contains numeric values", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 10),
      createMockCell("0_1", 20),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0", "0_1"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Average: 15")).toBeInTheDocument();
  });

  it("should not display Sum or Average when selection has no numeric values", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", "a"),
      createMockCell("0_1", "b"),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0", "0_1"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Count: 2")).toBeInTheDocument();
    expect(screen.queryByText(/Sum:/)).not.toBeInTheDocument();
    expect(screen.queryByText(/Average:/)).not.toBeInTheDocument();
  });

  it("should round sum and average to 3 decimal places", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 0.112_233_441_1),
      createMockCell("0_1", 0.112_233_441_1),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0", "0_1"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Sum: 0.224")).toBeInTheDocument();
    expect(screen.getByText("Average: 0.112")).toBeInTheDocument();
  });

  it("should correctly round sum and average to 3 decimal places", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 0.112_833_443_3),
      createMockCell("0_1", 0.112_833_443_3),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0", "0_1"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Sum: 0.226")).toBeInTheDocument();
    expect(screen.getByText("Average: 0.113")).toBeInTheDocument();
  });

  it("should not add extra decimal places to sum and average", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 0.1),
      createMockCell("0_1", 0.2),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0", "0_1"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Sum: 0.3")).toBeInTheDocument();
    expect(screen.getByText("Average: 0.15")).toBeInTheDocument();
  });

  it("should display Sum: 0 and Average: 0 when all selected values are zero", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 0),
      createMockCell("0_1", 0),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0", "0_1"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Count: 2")).toBeInTheDocument();
    expect(screen.getByText("Sum: 0")).toBeInTheDocument();
    expect(screen.getByText("Average: 0")).toBeInTheDocument();
  });

  it("should show hint instead of stats when exactly one cell is selected", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 10),
      createMockCell("0_1", 20),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("No selection")).toBeInTheDocument();
    expect(screen.queryByText(/Count:/)).not.toBeInTheDocument();
  });

  it("should display Sum and Average for negative numbers", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", -10),
      createMockCell("0_1", -20),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0", "0_1"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Count: 2")).toBeInTheDocument();
    expect(screen.getByText("Sum: -30")).toBeInTheDocument();
    expect(screen.getByText("Average: -15")).toBeInTheDocument();
  });

  it("should display Sum and Average for number strings", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", "42"),
      createMockCell("0_1", "3.14"),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness table={table} selectedCellIds={new Set(["0_0", "0_1"])} />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Count: 2")).toBeInTheDocument();
    expect(screen.getByText("Sum: 45.14")).toBeInTheDocument();
    expect(screen.getByText("Average: 22.57")).toBeInTheDocument();
  });

  it("should skip NaN and Infinity and show stats for finite values only", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 5),
      createMockCell("0_1", Number.NaN),
      createMockCell("0_2", Infinity),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness
          table={table}
          selectedCellIds={new Set(["0_0", "0_1", "0_2"])}
        />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Count: 3")).toBeInTheDocument();
    expect(screen.getByText("Sum: 5")).toBeInTheDocument();
    expect(screen.getByText("Average: 5")).toBeInTheDocument();
  });

  it("should exclude booleans from Sum and Average", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 10),
      createMockCell("0_1", true),
      createMockCell("0_2", false),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness
          table={table}
          selectedCellIds={new Set(["0_0", "0_1", "0_2"])}
        />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Count: 3")).toBeInTheDocument();
    expect(screen.getByText("Sum: 10")).toBeInTheDocument();
    expect(screen.getByText("Average: 10")).toBeInTheDocument();
  });

  it("should ignore select checkbox column for Count, Sum and Average", () => {
    const selectCellId = `0_${SELECT_COLUMN_ID}`;
    const row = createMockRow("0", [
      createMockCell(selectCellId, 10),
      createMockCell("0_0", 20),
      createMockCell("0_1", 30),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness
          table={table}
          selectedCellIds={new Set([selectCellId, "0_0", "0_1"])}
        />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Count: 2")).toBeInTheDocument();
    expect(screen.getByText("Sum: 50")).toBeInTheDocument();
    expect(screen.getByText("Average: 25")).toBeInTheDocument();
  });

  it("should show hint when only checkbox column cells are selected", () => {
    const selectCellId1 = `0_${SELECT_COLUMN_ID}`;
    const selectCellId2 = `1_${SELECT_COLUMN_ID}`;
    const row1 = createMockRow("0", [createMockCell(selectCellId1, true)]);
    const row2 = createMockRow("1", [createMockCell(selectCellId2, false)]);
    const table = createMockTable([row1, row2], []);

    render(
      <CellSelectionProvider>
        <TestHarness
          table={table}
          selectedCellIds={new Set([selectCellId1, selectCellId2])}
        />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("No selection")).toBeInTheDocument();
    expect(screen.queryByText(/Count:/)).not.toBeInTheDocument();
  });

  it("should show Sum and Average only for numeric cells when selection has mixed types", () => {
    const row = createMockRow("0", [
      createMockCell("0_0", 10),
      createMockCell("0_1", "abc"),
      createMockCell("0_2", undefined),
    ]);
    const table = createMockTable([row], []);

    render(
      <CellSelectionProvider>
        <TestHarness
          table={table}
          selectedCellIds={new Set(["0_0", "0_1", "0_2"])}
        />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Count: 3")).toBeInTheDocument();
    expect(screen.getByText("Sum: 10")).toBeInTheDocument();
    expect(screen.getByText("Average: 10")).toBeInTheDocument();
  });

  it("should display stats for multiple numeric cells across rows", () => {
    const row1 = createMockRow("0", [
      createMockCell("0_0", 1),
      createMockCell("0_1", 2),
    ]);
    const row2 = createMockRow("1", [
      createMockCell("1_0", 3),
      createMockCell("1_1", 4),
    ]);
    const table = createMockTable([row1, row2], []);

    render(
      <CellSelectionProvider>
        <TestHarness
          table={table}
          selectedCellIds={new Set(["0_0", "0_1", "1_0", "1_1"])}
        />
      </CellSelectionProvider>,
    );

    expect(screen.getByText("Count: 4")).toBeInTheDocument();
    expect(screen.getByText("Sum: 10")).toBeInTheDocument();
    expect(screen.getByText("Average: 2.5")).toBeInTheDocument();
  });
});