import { css } from "@emotion/react";
import type { ColumnDef, SortingFn } from "@tanstack/react-table";
import {
  flexRender,
  getCoreRowModel,
  getFilteredRowModel,
  getSortedRowModel,
  useReactTable,
} from "@tanstack/react-table";
import { useVirtualizer } from "@tanstack/react-virtual";
import { useMemo, useRef } from "react";
import { Focusable } from "react-aria";
import { graphql, useFragment } from "react-relay";

import {
  Flex,
  Heading,
  Icon,
  IconButton,
  Icons,
  Token,
  View,
} from "@phoenix/components";
import {
  RichTooltip,
  Tooltip,
  TooltipArrow,
  TooltipTrigger,
} from "@phoenix/components/core/tooltip";
import { Truncate } from "@phoenix/components/core/utility/Truncate";
import { GenerativeProviderIcon } from "@phoenix/components/generative/GenerativeProviderIcon";
import {
  getCommonPinningStyles,
  tableCSS,
} from "@phoenix/components/table/styles";
import { TimestampCell } from "@phoenix/components/table/TimestampCell";
import type {
  GenerativeModelKind,
  ModelsTable_generativeModels$data,
  ModelsTable_generativeModels$key,
} from "@phoenix/pages/settings/__generated__/ModelsTable_generativeModels.graphql";
import { EditModelButton } from "@phoenix/pages/settings/EditModelButton";
import type { Mutable } from "@phoenix/typeUtils";
import { getProviderName } from "@phoenix/utils/generativeUtils";
import { costFormatter } from "@phoenix/utils/numberFormatUtils";

import { CloneModelButton } from "./CloneModelButton";
import { DeleteModelButton } from "./DeleteModelButton";

// Initial estimate only — rows are measured once rendered. Roughly a 30px
// size-S icon button + 2 * 8px cell padding + 1px row border.
const ESTIMATED_ROW_HEIGHT = 47;

type ModelsTableGenerativeModel =
  ModelsTable_generativeModels$data["generativeModels"]["edges"][number]["generativeModel"];

type ModelsTableProps = {
  modelsRef: ModelsTable_generativeModels$key;
  kindFilter: "ALL" | GenerativeModelKind;
  search: string;
};

/**
 * Converts a date into a user search-able string, or undefined, so that tanstack table can filter on it
 * @param row - the row to get the date from
 * @returns the date, as a string, in the format of the DEFAULT_FORMAT
 */
function filterableDateAccessorFn(row?: string | null) {
  return row != null ? new Date(row).toISOString() : undefined;
}

/**
 * Gets the cost of a row for a given token type as a number, or undefined
 * @param row - the row to get the cost from
 * @param tokenType - the token type to get the cost for
 * @returns the cost of the row for the given token type
 */
function getRowCostNumber(row: ModelsTableGenerativeModel, tokenType: string) {
  const cost = row.tokenPrices?.find(
    (entry) => entry.tokenType === tokenType
  )?.costPerMillionTokens;
  return cost;
}

/**
 * Gets the cost of a row for a given token type as a string, or "--" if the cost is undefined
 * @param row - the row to get the cost from
 * @param tokenType - the token type to get the cost for
 * @returns the cost of the row for the given token type
 */
function getRowCost(row: ModelsTableGenerativeModel, tokenType: string) {
  const cost = getRowCostNumber(row, tokenType);
  return cost != null ? `${costFormatter(cost)}` : "--";
}

/**
 * Sorts a row by the cost of a given token type, ignoring the user-visible cost string during sort
 *
 * undefined values will float to the bottom.
 *
 * @param rowA - the first row to sort
 * @param rowB - the second row to sort
 * @param columnId - the id of the column to sort by
 * @returns the difference between the costs of the two rows
 */
const sortCostColumnFn: SortingFn<ModelsTableGenerativeModel> = (
  rowA,
  rowB,
  columnId
) => {
  const costA = getRowCostNumber(rowA.original, columnId);
  const costB = getRowCostNumber(rowB.original, columnId);
  if (costA == null || costB == null) {
    return 0;
  }
  return costA - costB;
};

export function ModelsTable({
  modelsRef,
  kindFilter,
  search,
}: ModelsTableProps) {
  "use no memo";
  const data = useFragment(
    graphql`
      fragment ModelsTable_generativeModels on Query
      @refetchable(queryName: "ModelsTableGenerativeModelsQuery")
      @argumentDefinitions(
        after: { type: "String", defaultValue: null }
        first: { type: "Int", defaultValue: 1000 }
      ) {
        generativeModels(first: $first, after: $after)
          @connection(key: "ModelsTable_generativeModels") {
          edges {
            generativeModel: node {
              id
              name
              provider
              namePattern
              providerKey
              startTime
              createdAt
              updatedAt
              lastUsedAt
              kind
              tokenPrices {
                tokenType
                kind
                costPerMillionTokens
              }
            }
          }
        }
      }
    `,
    modelsRef
  );

  const tableData = useMemo(
    () =>
      (data.generativeModels.edges.map((edge) => edge.generativeModel) ??
        []) as Mutable<
        (typeof data.generativeModels.edges)[number]["generativeModel"]
      >[],
    [data]
  );

  type TableRow = (typeof tableData)[number];
  const columns = useMemo(() => {
    const cols = [
      // invisible column for filtering purposes
      {
        id: "kind",
        enableHiding: true,
        accessorFn: (row) => row.kind,
      },
      {
        header: "name",
        accessorKey: "name",
        size: 200,
        cell: ({ row }) => {
          const model = row.original;
          return (
            <Flex
              direction="row"
              gap="size-100"
              alignItems="center"
              justifyContent="space-between"
            >
              <Truncate maxWidth="100%" title={model.name}>
                {model.name}
              </Truncate>
              <View flex="none">
                {row.original.kind === "CUSTOM" ? (
                  <Token>custom</Token>
                ) : (
                  <Token color="var(--global-color-blue-500)">built-in</Token>
                )}
              </View>
            </Flex>
          );
        },
      },
      {
        header: "provider",
        accessorFn: (row) => row.provider ?? undefined,
        sortUndefined: "last",
        cell: ({ row }) => {
          const providerKey = row.original.providerKey;
          if (!providerKey) {
            return <span>{row.original.provider || "--"}</span>;
          }
          return (
            <Flex direction="row" gap="size-100" alignItems="center">
              <GenerativeProviderIcon provider={providerKey} height={18} />
              <span>{getProviderName(providerKey as ModelProvider)}</span>
            </Flex>
          );
        },
      },
      {
        header: "match pattern",
        accessorKey: "namePattern",
        size: 200,
        cell: ({ row }) => {
          const namePattern = row.original.namePattern;
          const provider = row.original.provider;
          const providerKey = row.original.providerKey;
          return (
            <Flex direction="row" gap="size-100" alignItems="center">
              <Truncate maxWidth="100%" title={namePattern}>
                {namePattern}
              </Truncate>
              {provider && (
                <TooltipTrigger delay={0}>
                  <Focusable>
                    <Token role="button" color="var(--global-color-gray-300)">
                      <Flex direction="row" gap="size-100" alignItems="center">
                        {providerKey ? (
                          <>
                            <GenerativeProviderIcon
                              provider={providerKey}
                              height={18}
                            />
                            <span>{getProviderName(providerKey)}</span>
                          </>
                        ) : (
                          <span>{provider}</span>
                        )}
                      </Flex>
                    </Token>
                  </Focusable>
                  <Tooltip>
                    <TooltipArrow />
                    The preceding name pattern match will only be applied on
                    incoming traces with &quot;{provider}&quot; set as the
                    provider.
                  </Tooltip>
                </TooltipTrigger>
              )}
            </Flex>
          );
        },
      },
      {
        header: "prompt cost",
        id: "input",
        accessorFn: (row) => getRowCost(row, "input"),
        meta: { textAlign: "right" },
        sortingFn: sortCostColumnFn,
        sortUndefined: "last",
        cell: ({ row }) => {
          const inputCost = getRowCost(row.original, "input");
          return (
            <Flex justifyContent="end" alignItems="center" gap="size-100">
              <span>{inputCost}</span>
              <TooltipTrigger delay={0}>
                <IconButton size="S" aria-label="Input cost details">
                  <Icon svg={<Icons.MoreHorizontal />} />
                </IconButton>
                <RichTooltip>
                  <TooltipArrow />
                  <Heading level={4} weight="heavy">
                    Input Cost Breakdown
                  </Heading>
                  <View paddingTop="size-100">
                    {row.original.tokenPrices
                      ?.filter((entry) => entry.kind === "PROMPT")
                      .map((entry) => (
                        <Flex
                          key={entry.tokenType}
                          direction="row"
                          gap="size-100"
                          justifyContent="space-between"
                          width="100%"
                        >
                          <span>{entry.tokenType.toLocaleLowerCase()}</span>
                          <span>
                            {entry.costPerMillionTokens != null
                              ? `${costFormatter(entry.costPerMillionTokens)} / 1M`
                              : "--"}
                          </span>
                        </Flex>
                      )) ?? null}
                  </View>
                </RichTooltip>
              </TooltipTrigger>
            </Flex>
          );
        },
      },
      {
        header: "completion cost",
        id: "output",
        accessorFn: (row) => getRowCost(row, "output"),
        meta: { textAlign: "right" },
        sortingFn: sortCostColumnFn,
        sortUndefined: "last",
        cell: ({ row }) => {
          const outputCost = getRowCost(row.original, "output");
          return (
            <Flex justifyContent="end" alignItems="center" gap="size-100">
              <span>{outputCost}</span>
              <TooltipTrigger delay={0}>
                <IconButton size="S" aria-label="Output cost details">
                  <Icon svg={<Icons.MoreHorizontal />} />
                </IconButton>
                <RichTooltip>
                  <TooltipArrow />
                  <Heading level={4} weight="heavy">
                    Output Cost Breakdown
                  </Heading>
                  <View paddingTop="size-100">
                    {row.original.tokenPrices
                      ?.filter((entry) => entry.kind === "COMPLETION")
                      .map((entry) => (
                        <Flex
                          key={entry.tokenType}
                          direction="row"
                          gap="size-100"
                          justifyContent="space-between"
                          width="100%"
                        >
                          <span>{entry.tokenType.toLocaleLowerCase()}</span>
                          <span>
                            {entry.costPerMillionTokens != null
                              ? `${costFormatter(entry.costPerMillionTokens)} / 1M`
                              : "--"}
                          </span>
                        </Flex>
                      )) ?? null}
                  </View>
                </RichTooltip>
              </TooltipTrigger>
            </Flex>
          );
        },
      },
      {
        header: "start date",
        sortUndefined: "last",
        accessorFn: (row) => filterableDateAccessorFn(row.startTime),
        cell: TimestampCell,
      },
      {
        header: "created at",
        accessorFn: (row) => filterableDateAccessorFn(row.createdAt),
        sortUndefined: "last",
        cell: TimestampCell,
      },
      {
        header: "updated at",
        accessorFn: (row) => filterableDateAccessorFn(row.updatedAt),
        sortUndefined: "last",
        cell: TimestampCell,
      },
      {
        header: "last used at",
        accessorFn: (row) => filterableDateAccessorFn(row.lastUsedAt),
        sortUndefined: "last",
        cell: TimestampCell,
      },
      {
        id: "actions",
        header: "",
        accessorKey: "id",
        size: 120,
        enableResizing: false,
        cell: ({ row }) => {
          const isCustomModel = row.original.kind === "CUSTOM";
          return (
            <Flex
              direction="row"
              gap="size-50"
              width="100%"
              justifyContent="end"
            >
              {isCustomModel && (
                <TooltipTrigger>
                  <EditModelButton modelId={row.original.id} />
                  <Tooltip>
                    <TooltipArrow />
                    Edit model
                  </Tooltip>
                </TooltipTrigger>
              )}
              <TooltipTrigger>
                <CloneModelButton modelId={row.original.id} />
                <Tooltip>
                  <TooltipArrow />
                  Clone model
                </Tooltip>
              </TooltipTrigger>
              {isCustomModel && (
                <TooltipTrigger>
                  <DeleteModelButton
                    modelId={row.original.id}
                    modelName={row.original.name}
                  />
                  <Tooltip>
                    <TooltipArrow />
                    Delete model
                  </Tooltip>
                </TooltipTrigger>
              )}
            </Flex>
          );
        },
      },
    ] satisfies ColumnDef<TableRow>[];
    return cols;
  }, []);

  // if this is not memoized, the table will get into an infinite loop
  const columnFilters = useMemo(() => {
    return [
      {
        id: "kind",
        value: kindFilter === "ALL" ? "" : kindFilter,
      },
    ];
  }, [kindFilter]);

  // eslint-disable-next-line react-hooks-js/incompatible-library
  const table = useReactTable({
    columns,
    data: tableData,
    defaultColumn: {
      minSize: 100,
    },
    columnResizeMode: "onChange",
    state: {
      columnVisibility: {
        kind: false,
        provider: false,
      },
      columnPinning: {
        left: ["name"],
        right: ["actions"],
      },
      globalFilter: search?.trim(),
      columnFilters,
    },
    getCoreRowModel: getCoreRowModel(),
    getSortedRowModel: getSortedRowModel(),
    getFilteredRowModel: getFilteredRowModel(),
  });

  const rows = table.getRowModel().rows;
  const isEmpty = rows.length === 0;

  const tableContainerRef = useRef<HTMLDivElement>(null);
  // eslint-disable-next-line react-hooks-js/incompatible-library
  const virtualizer = useVirtualizer({
    count: rows.length,
    getScrollElement: () => tableContainerRef.current,
    estimateSize: () => ESTIMATED_ROW_HEIGHT,
    // Rows have fractional heights; the default integer offsetHeight
    // accumulates enough error over hundreds of rows to clip the last one.
    measureElement: (element) => element.getBoundingClientRect().height,
    overscan: 5,
  });
  const virtualRows = virtualizer.getVirtualItems();
  const totalHeight = virtualizer.getTotalSize();
  // Spacer rows keep the scrollbar sized to the full row count while only the
  // visible rows are mounted. Rows are not transformed so the pinned (sticky)
  // name and actions columns keep working.
  const paddingTop = virtualRows.length > 0 ? virtualRows[0].start : 0;
  const paddingBottom =
    virtualRows.length > 0
      ? totalHeight - virtualRows[virtualRows.length - 1].end
      : 0;

  if (isEmpty) {
    return (
      <Flex width="100%" justifyContent="center" alignItems="center">
        <p>No models found</p>
      </Flex>
    );
  }

  return (
    <div
      ref={tableContainerRef}
      css={css`
        // Fill the card body exactly so this div (the virtualizer's scroll
        // element) is the one that scrolls, keeping the mounted rows in sync
        // with the viewport.
        height: 100%;
        box-sizing: border-box;
        overflow: auto;
      `}
    >
      <table
        css={[
          tableCSS,
          // Fixed layout so long model names and match patterns truncate at
          // their column size instead of stretching the column.
          css`
            table-layout: fixed;
          `,
        ]}
        style={{ width: table.getTotalSize(), minWidth: "100%" }}
      >
        <thead>
          {table.getHeaderGroups().map((headerGroup) => (
            <tr key={headerGroup.id}>
              {headerGroup.headers.map((header) => (
                <th
                  colSpan={header.colSpan}
                  key={header.id}
                  style={{
                    ...getCommonPinningStyles(header.column),
                    boxSizing: "border-box",
                    minWidth: header.column.getSize(),
                    maxWidth: header.column.getSize(),
                  }}
                >
                  {header.isPlaceholder ? null : (
                    <>
                      <div
                        {...{
                          className: header.column.getCanSort() ? "sort" : "",
                          onClick: header.column.getToggleSortingHandler(),
                          style: {
                            textWrap: "nowrap",
                            textAlign: header.column.columnDef.meta?.textAlign,
                          },
                        }}
                      >
                        {flexRender(
                          header.column.columnDef.header,
                          header.getContext()
                        )}
                        {header.column.getIsSorted() ? (
                          <Icon
                            className="sort-icon"
                            svg={
                              header.column.getIsSorted() === "asc" ? (
                                <Icons.CaretUpFilled />
                              ) : (
                                <Icons.CaretDownFilled />
                              )
                            }
                          />
                        ) : null}
                      </div>
                      {header.column.getCanResize() && (
                        <div
                          {...{
                            onMouseDown: header.getResizeHandler(),
                            onTouchStart: header.getResizeHandler(),
                            className: `resizer ${
                              header.column.getIsResizing() ? "isResizing" : ""
                            }`,
                          }}
                        />
                      )}
                    </>
                  )}
                </th>
              ))}
            </tr>
          ))}
        </thead>
        <tbody>
          {paddingTop > 0 && (
            <tr>
              <td
                colSpan={table.getVisibleLeafColumns().length}
                style={{ height: paddingTop, padding: 0, border: "none" }}
              />
            </tr>
          )}
          {virtualRows.map((virtualRow) => {
            const row = rows[virtualRow.index];
            return (
              <tr
                key={row.id}
                data-index={virtualRow.index}
                ref={virtualizer.measureElement}
              >
                {row.getVisibleCells().map((cell) => (
                  <td
                    key={cell.id}
                    align={cell.column.columnDef.meta?.textAlign}
                    style={{
                      ...getCommonPinningStyles(cell.column),
                    }}
                  >
                    {flexRender(cell.column.columnDef.cell, cell.getContext())}
                  </td>
                ))}
              </tr>
            );
          })}
          {paddingBottom > 0 && (
            <tr>
              <td
                colSpan={table.getVisibleLeafColumns().length}
                style={{ height: paddingBottom, padding: 0, border: "none" }}
              />
            </tr>
          )}
        </tbody>
      </table>
    </div>
  );
}
