"use client"

import * as React from "react"
import { type Column, type Table } from "@tanstack/react-table"
import { Check, ChevronsUpDown, MoveRight } from "lucide-react"

import { cn, toSentenceCase } from "@/lib/utils"
import { Button } from "@/components/ui/button"
import {
  Command,
  CommandEmpty,
  CommandGroup,
  CommandInput,
  CommandItem,
  CommandList,
  CommandSeparator,
} from "@/components/ui/command"
import {
  Popover,
  PopoverContent,
  PopoverTrigger,
} from "@/components/ui/popover"

/**
 * Helper function to check if a column is a parent column (has subcolumns)
 */
function isParentColumn<TData>(column: Column<TData>): boolean {
  return column.columns && column.columns.length > 0
}

/**
 * Helper function to pin all subcolumns of a parent column
 */
function pinSubColumns<TData>(
  column: Column<TData>,
  pinType: false | "left" | "right"
): void {
  // If this is a parent column, pin all its subcolumns
  if (isParentColumn(column)) {
    column.columns.forEach((subColumn) => {
      // Recursively handle nested columns
      pinSubColumns(subColumn, pinType)
    })
  } else {
    // For leaf columns, apply the pin if possible
    if (column.getCanPin?.()) {
      column.pin?.(pinType)
    }
  }
}

/**
 * Checks if all subcolumns of a parent column are pinned to the specified side
 */
function areAllSubColumnsPinned<TData>(
  column: Column<TData>,
  pinType: "left" | "right"
): boolean {
  if (isParentColumn(column)) {
    // Check if all subcolumns are pinned
    return column.columns.every((subColumn) => 
      areAllSubColumnsPinned(subColumn, pinType)
    )
  } else {
    // For leaf columns, check if it's pinned to the specified side
    return column.getIsPinned?.() === pinType
  }
}

/**
 * "Pin Right" Popover. Supports pinning both individual columns and header groups.
 */
export function PinRightButton<TData>({ table }: { table: Table<TData> }) {
  const [open, setOpen] = React.useState(false)
  const triggerRef = React.useRef<HTMLButtonElement>(null)

  // Get all columns that can be pinned, including parent columns
  const pinnableColumns = React.useMemo(() => {
    return table.getAllColumns().filter((column) => {
      // If it's a leaf column, check if it can be pinned
      if (!isParentColumn(column)) {
        return column.getCanPin?.()
      }
      
      // If it's a parent column, check if at least one subcolumn can be pinned
      return column.columns.some((subCol) => {
        if (isParentColumn(subCol)) {
          // Recursively check nested columns
          return subCol.columns.some(c => c.getCanPin?.())
        }
        return subCol.getCanPin?.()
      })
    })
  }, [table])

  // Handle column pinning
  const handleColumnPin = React.useCallback((column: Column<TData>) => {
    // For parent columns, pin/unpin all subcolumns
    if (isParentColumn(column)) {
      const allPinned = areAllSubColumnsPinned(column, "right")
      pinSubColumns(column, allPinned ? false : "right")
    } else {
      // For leaf columns, toggle pin state
      const isPinned = column.getIsPinned?.() === "right"
      column.pin?.(isPinned ? false : "right")
    }
  }, [])

  // Check if a column or its subcolumns are pinned right
  const isColumnPinned = React.useCallback((column: Column<TData>): boolean => {
    if (isParentColumn(column)) {
      return areAllSubColumnsPinned(column, "right")
    } else {
      return column.getIsPinned?.() === "right"
    }
  }, [])

  return (
    <Popover open={open} onOpenChange={setOpen}>
      <PopoverTrigger asChild>
        <Button
          ref={triggerRef}
          variant="outline"
          size="sm"
          className="h-8 gap-2"
        >
          <MoveRight className="size-4" />
          
          <span className="hidden sm:inline">
            Right
          </span>
          <ChevronsUpDown className="ml-1 size-4 opacity-50 hidden sm:inline" />
        </Button>
      </PopoverTrigger>
      
      <PopoverContent
        align="end"
        className="w-56 p-0"
        onCloseAutoFocus={() => triggerRef.current?.focus()}
      >
        <Command>
          <CommandInput placeholder="Search columns..." />
          <CommandList>
            <CommandEmpty>No columns found.</CommandEmpty>
            <CommandGroup>
              {/* Header Columns (Parent Columns) */}
              {pinnableColumns
                .filter(isParentColumn)
                .map((column) => (
                  <CommandItem
                    key={column.id}
                    onSelect={() => {
                      handleColumnPin(column)
                    }}
                    className="font-medium"
                  >
                    <span className="truncate">
                      {column.id === "Basic Info" || column.id === "Metadata" 
                        ? column.id // Use column ID directly for common groups
                        : toSentenceCase(column.id)}
                    </span>
                    <Check
                      className={cn(
                        "ml-auto size-4 shrink-0",
                        isColumnPinned(column) ? "opacity-100" : "opacity-0"
                      )}
                    />
                  </CommandItem>
                ))}
              
              {pinnableColumns.some(isParentColumn) && 
               pinnableColumns.some(col => !isParentColumn(col)) && (
                <CommandSeparator />
              )}
              
              {/* Leaf Columns (individual columns) */}
              {pinnableColumns
                .filter(col => !isParentColumn(col))
                .map((column) => (
                  <CommandItem
                    key={column.id}
                    onSelect={() => {
                      handleColumnPin(column)
                    }}
                  >
                    <span className="truncate">
                      {toSentenceCase(column.id)}
                    </span>
                    <Check
                      className={cn(
                        "ml-auto size-4 shrink-0",
                        isColumnPinned(column) ? "opacity-100" : "opacity-0"
                      )}
                    />
                  </CommandItem>
                ))}
            </CommandGroup>
          </CommandList>
        </Command>
      </PopoverContent>
    </Popover>
  )
}