"use client"

import * as React from "react"
import { ColumnDef } from "@tanstack/react-table"
import { Checkbox } from "@/components/ui/checkbox"
import { Button } from "@/components/ui/button"
import { Badge } from "@/components/ui/badge"
import { Ellipsis, Users, Boxes } from "lucide-react"
// import { toast } from "sonner" // If needed
import { DataTableColumnHeaderSimple } from "@/components/data-table/data-table-column-simple-header"
import { formatDate } from "@/lib/utils" // or your date util

// Example: If you have a type for row actions
import { type DataTableRowAction } from "@/types/table"
import { ContactItem, PossibleItem, vendorInvestigationsColumnsConfig, VendorInvestigationsViewWithContacts } from "@/config/vendorInvestigationsColumnsConfig"

// Props that define how we handle special columns (contacts, items, actions, etc.)
interface GetVendorInvestigationsColumnsProps {
  setRowAction?: React.Dispatch<
    React.SetStateAction<
      DataTableRowAction<VendorInvestigationsViewWithContacts> | null
    >
  >
  openContactsModal?: (investigationId: number, contacts: ContactItem[]) => void
  openItemsDrawer?: (investigationId: number, items: PossibleItem[]) => void
}

// This function returns the array of columns for TanStack Table
export function getColumns({
  setRowAction,
  openContactsModal,
  openItemsDrawer,
}: GetVendorInvestigationsColumnsProps): ColumnDef<
  VendorInvestigationsViewWithContacts
>[] {
  // --------------------------------------------
  // 1) Select (checkbox) column
  // --------------------------------------------
  const selectColumn: ColumnDef<VendorInvestigationsViewWithContacts> = {
    id: "select",
    header: ({ table }) => (
      <Checkbox
        checked={
          table.getIsAllPageRowsSelected() ||
          (table.getIsSomePageRowsSelected() && "indeterminate")
        }
        onCheckedChange={(value) => table.toggleAllPageRowsSelected(!!value)}
        aria-label="Select all"
        className="translate-y-0.5"
      />
    ),
    cell: ({ row }) => (
      <Checkbox
        checked={row.getIsSelected()}
        onCheckedChange={(value) => row.toggleSelected(!!value)}
        aria-label="Select row"
        className="translate-y-0.5"
      />
    ),
    size: 40,
    enableSorting: false,
    enableHiding: false,
  }

  // --------------------------------------------
  // 2) Actions column (optional)
  // --------------------------------------------
  const actionsColumn: ColumnDef<VendorInvestigationsViewWithContacts> = {
    id: "actions",
    enableHiding: false,
    cell: ({ row }) => {
      const inv = row.original

      return (
        <Button
          variant="ghost"
          className="flex size-8 p-0 data-[state=open]:bg-muted"
          aria-label="Open menu"
          onClick={() => {
            // e.g. open a dropdown or set your row action
            setRowAction?.({ type: "update", row })
          }}
        >
          <Ellipsis className="size-4" aria-hidden="true" />
        </Button>
      )
    },
    size: 40,
  }

  // --------------------------------------------
  // 3) Contacts column (badge count -> open modal)
  // --------------------------------------------
  const contactsColumn: ColumnDef<VendorInvestigationsViewWithContacts> = {
    id: "contacts",
    header: "Contacts",
    cell: ({ row }) => {
      const { contacts, investigationId } = row.original
      const count = contacts?.length ?? 0

      const handleClick = () => {
        openContactsModal?.(investigationId, contacts)
      }

      return (
        <Button
          variant="ghost"
          size="sm"
          className="relative h-8 w-8 p-0 group"
          onClick={handleClick}
          aria-label={
            count > 0 ? `View ${count} contacts` : "Add contacts"
          }
        >
          <Users className="h-4 w-4 text-muted-foreground group-hover:text-primary transition-colors" />
          {count > 0 && (
            <Badge
              variant="secondary"
              className="pointer-events-none absolute -top-1 -right-1 h-4 min-w-[1rem] p-0 text-[0.625rem] leading-none flex items-center justify-center"
            >
              {count}
            </Badge>
          )}
          <span className="sr-only">
            {count > 0 ? `${count} Contacts` : "Add Contacts"}
          </span>
        </Button>
      )
    },
    enableSorting: false,
    size: 60,
  }

  // --------------------------------------------
  // 4) Possible Items column (badge count -> open drawer)
  // --------------------------------------------
  const possibleItemsColumn: ColumnDef<VendorInvestigationsViewWithContacts> = {
    id: "possibleItems",
    header: "Items",
    cell: ({ row }) => {
      const { possibleItems, investigationId } = row.original
      const count = possibleItems?.length ?? 0

      const handleClick = () => {
        openItemsDrawer?.(investigationId, possibleItems)
      }

      return (
        <Button
          variant="ghost"
          size="sm"
          className="relative h-8 w-8 p-0 group"
          onClick={handleClick}
          aria-label={
            count > 0 ? `View ${count} items` : "Add items"
          }
        >
          <Boxes className="h-4 w-4 text-muted-foreground group-hover:text-primary transition-colors" />
          {count > 0 && (
            <Badge
              variant="secondary"
              className="pointer-events-none absolute -top-1 -right-1 h-4 min-w-[1rem] p-0 text-[0.625rem] leading-none flex items-center justify-center"
            >
              {count}
            </Badge>
          )}
          <span className="sr-only">
            {count > 0 ? `${count} Items` : "Add Items"}
          </span>
        </Button>
      )
    },
    enableSorting: false,
    size: 60,
  }

  // --------------------------------------------
  // 5) Build "grouped" columns from config
  // --------------------------------------------
  const groupMap: Record<string, ColumnDef<VendorInvestigationsViewWithContacts>[]> = {}

  vendorInvestigationsColumnsConfig.forEach((cfg) => {
    const groupName = cfg.group || "_noGroup"

    if (!groupMap[groupName]) {
      groupMap[groupName] = []
    }

    const childCol: ColumnDef<VendorInvestigationsViewWithContacts> = {
      accessorKey: cfg.id,
      header: ({ column }) => (
        <DataTableColumnHeaderSimple column={column} title={cfg.label} />
      ),
      meta: {
        excelHeader: cfg.excelHeader,
        group: cfg.group,
        type: cfg.type,
      },
      cell: ({ row, cell }) => {
        const val = cell.getValue()

        // Example: Format date fields
        if (
          cfg.id === "investigationCreatedAt" ||
          cfg.id === "investigationUpdatedAt" ||
          cfg.id === "scheduledStartAt" ||
          cfg.id === "scheduledEndAt" ||
          cfg.id === "completedAt"
        ) {
          const dateVal = val ? new Date(val as string) : null
          return dateVal ? formatDate(dateVal) : ""
        }

        // Example: You could show an icon for "investigationStatus"
        if (cfg.id === "investigationStatus") {
          return <span className="capitalize">{val as string}</span>
        }

        return val ?? ""
      },
    }

    groupMap[groupName].push(childCol)
  })

  // Turn the groupMap into nested columns
  const nestedColumns: ColumnDef<VendorInvestigationsViewWithContacts>[] = []
  for (const [groupName, colDefs] of Object.entries(groupMap)) {
    if (groupName === "_noGroup") {
      nestedColumns.push(...colDefs)
    } else {
      nestedColumns.push({
        id: groupName,
        header: groupName,
        columns: colDefs,
      })
    }
  }

  // --------------------------------------------
  // 6) Return final columns array
  //    (You can reorder these as you wish.)
  // --------------------------------------------
  return [
    selectColumn,
    ...nestedColumns,
    contactsColumn,
    possibleItemsColumn,
    actionsColumn,
  ]
}