import React, {
  createContext,
  useContext,
  useLayoutEffect,
  useMemo,
} from "react";
import { useTableActions, useTableStore } from "./table";

interface TheadContextValue {
  columnCount: number;
  columnWidths: number[];
  // Context exists to ensure Th is wrapped in Thead
  // The colIndex is injected via props by React.cloneElement
}

const TheadContext = createContext<TheadContextValue | null>(null);

export const useTheadContext = (): TheadContextValue => {
  const context = useContext(TheadContext);
  if (!context) {
    throw new Error("Th component must be used inside Thead component");
  }
  return context;
};

interface IProps extends React.HTMLAttributes<HTMLDivElement> {
  headerHeight?: number;
}
const Thead = ({ children, style, headerHeight = 50, ...props }: IProps) => {
  const contentWidth = useTableStore((s) => s.contentWidth);
  const columnCount = useTableStore((s) => s.columnCount);
  const columnWidths = useTableStore((s) => s.columnWidths);
  const { setColumnWidths } = useTableActions();

  // Extract widths from Th children and update table context
  // The setColumnWidths function now handles change detection internally
  useLayoutEffect(() => {
    const widths: number[] = [];
    React.Children.forEach(children, (child) => {
      if (React.isValidElement(child)) {
        const thProps = child.props as { width?: number; colSpan?: number };
        const colSpan = thProps.colSpan || 1;
        const totalWidth = thProps.width ?? 100;
        const widthPerCol = totalWidth / colSpan;

        for (let i = 0; i < colSpan; i++) {
          widths.push(widthPerCol);
        }
      }
    });
    if (widths.length > 0) {
      setColumnWidths(widths);
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [children, setColumnWidths]);

  const contextValue: TheadContextValue = useMemo(
    () => ({
      columnCount,
      columnWidths,
    }),
    [columnCount, columnWidths],
  );

  let currentColIndex = 0;

  return (
    <TheadContext.Provider value={contextValue}>
      <div
        style={{
          position: "sticky",
          top: 0,
          zIndex: 10,
          display: "flex",
          height: headerHeight,
          width: "100%",
          minWidth: contentWidth,
          boxSizing: "border-box",
          ...style,
          backgroundColor: "#f5f5f5",
          borderBottom: "2px solid #ddd",
        }}
        {...props}
      >
        {React.Children.map(children, (child) => {
          if (React.isValidElement(child)) {
            const colSpan = (child.props as any).colSpan || 1;
            const index = currentColIndex;
            currentColIndex += colSpan;
            return React.cloneElement(child, { colIndex: index } as any);
          }
          return child;
        })}
        {/* Spacer to fill remaining space when content is smaller than container */}
        {/* {needsFill && spacerWidth > 0 && (
          <div
            style={{
              width: spacerWidth,
              minWidth: spacerWidth,
              maxWidth: spacerWidth,
              flexShrink: 0,
              flexGrow: 0,
            }}
          />
        )} */}
      </div>
    </TheadContext.Provider>
  );
};

export default Thead;