import { FC, useState, useEffect } from "react";
import {
    Box, Checkbox, CheckboxGroup, VStack, HStack, Text, Button,
    Accordion, AccordionItem, AccordionButton, AccordionPanel, AccordionIcon,
    Badge
} from "@chakra-ui/react";

export type SheetColumnInfo = {
    sheetName: string;
    columns: string[];
};

export type ColumnSelection = {
    sheetName: string;
    columnIndices: number[];
};

type Props = {
    sheetColumns: SheetColumnInfo[];
    onChange: (selections: ColumnSelection[]) => void;
};

const ColumnSelector: FC<Props> = ({ sheetColumns, onChange }) => {
    const [selections, setSelections] = useState<Record<string, Set<number>>>(() => {
        const init: Record<string, Set<number>> = {};
        sheetColumns.forEach(s => {
            init[s.sheetName] = new Set(s.columns.map((_, i) => i));
        });
        return init;
    });

    useEffect(() => {
        const result: ColumnSelection[] = sheetColumns.map(s => ({
            sheetName: s.sheetName,
            columnIndices: Array.from(selections[s.sheetName] ?? [])
        }));
        onChange(result);
    }, [selections, sheetColumns, onChange]);

    const toggleColumn = (sheetName: string, colIdx: number) => {
        setSelections(prev => {
            const set = new Set(prev[sheetName] ?? []);
            if (set.has(colIdx)) set.delete(colIdx);
            else set.add(colIdx);
            return { ...prev, [sheetName]: set };
        });
    };

    const selectAll = (sheetName: string, cols: string[]) => {
        setSelections(prev => ({
            ...prev,
            [sheetName]: new Set(cols.map((_, i) => i))
        }));
    };

    const selectNone = (sheetName: string) => {
        setSelections(prev => ({ ...prev, [sheetName]: new Set() }));
    };

    return (
        <Box w="full" borderWidth={1} borderRadius="md" p={3}>
            <Text fontWeight="bold" mb={2} fontSize="sm">Select Columns to Translate</Text>
            <Accordion allowMultiple defaultIndex={sheetColumns.map((_, i) => i)}>
                {sheetColumns.map(sheet => (
                    <AccordionItem key={sheet.sheetName}>
                        <AccordionButton>
                            <Box flex="1" textAlign="left" fontSize="sm" fontWeight="semibold">
                                {sheet.sheetName}
                                <Badge ml={2} colorScheme="lingva">
                                    {selections[sheet.sheetName]?.size ?? 0}/{sheet.columns.length}
                                </Badge>
                            </Box>
                            <AccordionIcon />
                        </AccordionButton>
                        <AccordionPanel pb={3}>
                            <HStack mb={2} spacing={2}>
                                <Button size="xs" variant="outline" colorScheme="lingva"
                                    onClick={() => selectAll(sheet.sheetName, sheet.columns)}>
                                    All
                                </Button>
                                <Button size="xs" variant="outline"
                                    onClick={() => selectNone(sheet.sheetName)}>
                                    None
                                </Button>
                            </HStack>
                            <VStack align="start" spacing={1} maxH="150px" overflowY="auto">
                                {sheet.columns.map((col, idx) => (
                                    <Checkbox
                                        key={idx}
                                        size="sm"
                                        isChecked={selections[sheet.sheetName]?.has(idx) ?? false}
                                        onChange={() => toggleColumn(sheet.sheetName, idx)}
                                        colorScheme="lingva"
                                    >
                                        <Text fontSize="xs">{col}</Text>
                                    </Checkbox>
                                ))}
                            </VStack>
                        </AccordionPanel>
                    </AccordionItem>
                ))}
            </Accordion>
        </Box>
    );
};

export default ColumnSelector;