import React, { useState, useEffect, useImperativeHandle, forwardRef } from 'react';
import { Checkbox } from '@mui/material';

const SelectionHeader = forwardRef((params, ref) => {
  const [checkboxState, setCheckboxState] = useState({ checked: false, indeterminate: false });
  
  // Expose updateState method to parent via ref
  useImperativeHandle(ref, () => ({
    updateState: (selectedRows, displayedRows) => {
      const allSelected = displayedRows > 0 && selectedRows.size >= displayedRows;
      const someSelected = selectedRows.size > 0;
      const indeterminate = someSelected && !allSelected;
      setCheckboxState({ checked: allSelected, indeterminate });
    }
  }));
  
  // Register this component with parent on mount
  useEffect(() => {
    if (params.api?.setHeaderCheckboxRef) {
      params.api.setHeaderCheckboxRef(ref.current || { updateState: (selectedRows, displayedRows) => {
        const allSelected = displayedRows > 0 && selectedRows.size >= displayedRows;
        const someSelected = selectedRows.size > 0;
        const indeterminate = someSelected && !allSelected;
        setCheckboxState({ checked: allSelected, indeterminate });
      }});
    }
  }, [params.api, ref]);
  
  return (
    <div 
      style={{ 
        display: 'flex', 
        alignItems: 'center', 
        justifyContent: 'center',
        height: '100%',
        width: '100%'
      }}
    >
      <Checkbox
        size="small"
        checked={checkboxState.checked}
        indeterminate={checkboxState.indeterminate}
        onClick={(event) => {
          event.stopPropagation();
          
          const parentComponent = params.api?.fibdashComponent;
          if (parentComponent && parentComponent.onSelectAll) {
            const { selectedRows } = parentComponent.state;
            const shouldSelectAll = selectedRows.size === 0;
            parentComponent.onSelectAll(shouldSelectAll);
          }
        }}
        sx={{
          padding: 0,
          '& .MuiSvgIcon-root': {
            fontSize: 18
          }
        }}
      />
    </div>
  );
});

export default SelectionHeader;