import { Children, isValidElement, useMemo, useState, type ReactNode } from "react";
import { ArrowDown, ArrowUp, ArrowUpDown, RotateCcw } from "lucide-react";

type SortDirection = "asc" | "desc";

interface Cell {
  content: ReactNode;
  text: string;
}

function extractText(node: ReactNode): string {
  if (node === null || node === undefined || typeof node === "boolean") return "";
  if (typeof node === "string" || typeof node === "number") return String(node);
  if (Array.isArray(node)) return node.map(extractText).join("");
  if (isValidElement(node)) return extractText((node.props as { children?: ReactNode }).children);
  return "";
}

function toCells(row: ReactNode): Cell[] {
  return Children.toArray(row)
    .filter(isValidElement)
    .map((cell) => {
      const content = (cell.props as { children?: ReactNode }).children;
      return { content, text: extractText(content).trim() };
    });
}

function toRows(section: ReactNode): Cell[][] {
  const rows = Children.toArray(section).filter(isValidElement);
  return rows.map((row) => toCells((row.props as { children?: ReactNode }).children));
}

function compareCells(a: Cell, b: Cell, direction: SortDirection) {
  const numA = Number(a.text.replace(/[,$%\s]/g, ""));
  const numB = Number(b.text.replace(/[,$%\s]/g, ""));
  const bothNumeric = a.text !== "" && b.text !== "" && !Number.isNaN(numA) && !Number.isNaN(numB);
  const result = bothNumeric ? numA - numB : a.text.localeCompare(b.text, undefined, { numeric: true, sensitivity: "base" });
  return direction === "asc" ? result : -result;
}

export function SortableMarkdownTable({ children }: { children?: ReactNode }) {
  const sections = Children.toArray(children).filter(isValidElement);
  const theadEl = sections.find((section) => section.type === "thead");
  const tbodyEl = sections.find((section) => section.type === "tbody");

  const headerCells = theadEl ? toRows((theadEl.props as { children?: ReactNode }).children)[0] ?? [] : [];
  const bodyRows = tbodyEl ? toRows((tbodyEl.props as { children?: ReactNode }).children) : [];

  const [sortColumn, setSortColumn] = useState<number | null>(null);
  const [sortDirection, setSortDirection] = useState<SortDirection>("asc");

  const sortedRows = useMemo(() => {
    if (sortColumn === null) return bodyRows;
    return [...bodyRows].sort((rowA, rowB) => {
      const cellA = rowA[sortColumn];
      const cellB = rowB[sortColumn];
      if (!cellA || !cellB) return 0;
      return compareCells(cellA, cellB, sortDirection);
    });
  }, [bodyRows, sortColumn, sortDirection]);

  if (headerCells.length === 0) {
    // Not a well-formed GFM table (no thead) — fall back to plain rendering.
    return (
      <div className="my-4 overflow-x-auto">
        <table className="w-full table-auto border-collapse text-sm">{children}</table>
      </div>
    );
  }

  const toggleSort = (columnIndex: number) => {
    if (sortColumn !== columnIndex) {
      setSortColumn(columnIndex);
      setSortDirection("asc");
    } else if (sortDirection === "asc") {
      setSortDirection("desc");
    } else {
      setSortColumn(null);
    }
  };

  return (
    <div className="my-4">
      {sortColumn !== null && (
        <div className="mb-1.5 flex justify-end">
          <button
            type="button"
            onClick={() => setSortColumn(null)}
            className="inline-flex items-center gap-1 rounded-md px-2 py-1 text-xs font-medium text-slate-500 hover:bg-slate-100 hover:text-slate-800"
          >
            <RotateCcw className="h-3 w-3" />
            Reset sort
          </button>
        </div>
      )}
      <div className="overflow-x-auto">
        <table className="w-full table-auto border-collapse text-sm">
          <thead>
            <tr>
              {headerCells.map((cell, columnIndex) => {
                const isActive = sortColumn === columnIndex;
                const Icon = isActive ? (sortDirection === "asc" ? ArrowUp : ArrowDown) : ArrowUpDown;
                return (
                  <th
                    key={columnIndex}
                    onClick={() => toggleSort(columnIndex)}
                    aria-sort={isActive ? (sortDirection === "asc" ? "ascending" : "descending") : "none"}
                    className="cursor-pointer select-none border border-slate-200 bg-slate-50 px-3 py-2 text-left font-semibold text-slate-900 hover:bg-slate-100"
                  >
                    <span className="inline-flex items-center gap-1.5">
                      {cell.content}
                      <Icon className={isActive ? "h-3.5 w-3.5 text-emerald-700" : "h-3.5 w-3.5 text-slate-400"} />
                    </span>
                  </th>
                );
              })}
            </tr>
          </thead>
          <tbody>
            {sortedRows.map((row, rowIndex) => (
              <tr key={rowIndex}>
                {row.map((cell, columnIndex) => (
                  <td key={columnIndex} className="border border-slate-200 px-3 py-2 align-top text-slate-700">
                    {cell.content}
                  </td>
                ))}
              </tr>
            ))}
          </tbody>
        </table>
      </div>
    </div>
  );
}