import {
  Alert,
  Badge,
  Box,
  Button,
  Checkbox,
  Group,
  Loader,
  Paper,
  Stack,
  Table,
  Text,
} from '@mantine/core';
import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
import { ChevronDown, ChevronRight, TrainFront } from 'lucide-react';
import { useMemo, useState } from 'react';

import { useToast } from '@/hooks/use-toast';
import {
  warehouseService,
  type LoadableTrain,
  type TrainLoadableItem,
} from '@/services/warehouse.service';
import { extractErrorMessage } from './options';

const STAGE_COLOR: Record<string, string> = {
  RECEIVED: 'blue',
  STORED: 'gray',
  RESERVED: 'grape',
  READY_FOR_LOADING: 'teal',
  LOADED: 'green',
};

const weight = (w: number | null) => (w == null ? '—' : `${Number(w).toLocaleString()} t`);

interface BookingGroup {
  bookingId: string | null;
  bookingReference: string | null;
  customerName: string | null;
  items: TrainLoadableItem[];
}

function groupByBooking(items: TrainLoadableItem[]): BookingGroup[] {
  const map = new Map<string, BookingGroup>();
  for (const i of items) {
    const key = i.bookingId ?? i.bookingReference ?? 'unknown';
    let g = map.get(key);
    if (!g) {
      g = { bookingId: i.bookingId, bookingReference: i.bookingReference, customerName: i.customerName, items: [] };
      map.set(key, g);
    }
    g.items.push(i);
  }
  return [...map.values()];
}

/**
 * Load to Train — a datatable of allocated EXPORT trains. Expand a train to see
 * the bookings allocated to it; expand a booking to see its containers/cargoes
 * and load the ready ones onto their wagons. Only READY_FOR_LOADING items with an
 * allocated wagon are selectable.
 */
export function LoadToTrainPanel() {
  const { data: trains = [], isLoading } = useQuery({
    queryKey: ['loadable-trains'],
    queryFn: () => warehouseService.getLoadableTrains(),
  });
  const [expanded, setExpanded] = useState<Set<string>>(new Set());
  const toggle = (id: string) =>
    setExpanded((s) => {
      const next = new Set(s);
      next.has(id) ? next.delete(id) : next.add(id);
      return next;
    });

  if (isLoading) {
    return (
      <Group justify="center" py="lg">
        <Loader />
      </Group>
    );
  }
  if (trains.length === 0) {
    return (
      <Alert color="gray" variant="light">
        No allocated EXPORT trains awaiting loading. Trains appear here after train and wagon allocation.
      </Alert>
    );
  }

  return (
    <Table.ScrollContainer minWidth={720}>
      <Table verticalSpacing="sm">
        <Table.Thead>
          <Table.Tr>
            <Table.Th w={40} />
            <Table.Th>Train</Table.Th>
            <Table.Th>Route</Table.Th>
            <Table.Th ta="center">Ready</Table.Th>
            <Table.Th ta="center">Loaded</Table.Th>
          </Table.Tr>
        </Table.Thead>
        <Table.Tbody>
          {trains.map((t) => (
            <TrainRow
              key={t.scheduleId}
              train={t}
              expanded={expanded.has(t.scheduleId)}
              onToggle={() => toggle(t.scheduleId)}
            />
          ))}
        </Table.Tbody>
      </Table>
    </Table.ScrollContainer>
  );
}

function TrainRow({ train, expanded, onToggle }: { train: LoadableTrain; expanded: boolean; onToggle: () => void }) {
  const { data: items = [], isLoading } = useQuery({
    queryKey: ['train-loadable-items', train.scheduleId],
    queryFn: () => warehouseService.getTrainLoadableItems(train.scheduleId),
    enabled: expanded,
  });
  const bookings = useMemo(() => groupByBooking(items), [items]);
  const route =
    train.origin || train.destination ? `${train.origin ?? '?'} → ${train.destination ?? '?'}` : '—';

  return (
    <>
      <Table.Tr style={{ cursor: 'pointer' }} onClick={onToggle}>
        <Table.Td>{expanded ? <ChevronDown size={16} /> : <ChevronRight size={16} />}</Table.Td>
        <Table.Td>
          <Group gap="xs" wrap="nowrap">
            <TrainFront size={16} />
            <Text fw={600}>{train.trainNumber ?? train.scheduleId.slice(0, 8)}</Text>
          </Group>
        </Table.Td>
        <Table.Td>{route}</Table.Td>
        <Table.Td ta="center">
          <Badge color="blue" variant="light">
            {train.readyCount}
          </Badge>
        </Table.Td>
        <Table.Td ta="center">
          <Badge color="green" variant="light">
            {train.loadedCount}
          </Badge>
        </Table.Td>
      </Table.Tr>
      {expanded && (
        <Table.Tr>
          <Table.Td colSpan={5} p={0}>
            <Box p="sm" bg="var(--mantine-color-gray-0)">
              {isLoading ? (
                <Group justify="center" py="md">
                  <Loader size="sm" />
                </Group>
              ) : bookings.length === 0 ? (
                <Alert color="gray" variant="light">
                  No arrived containers/cargoes allocated to this train yet.
                </Alert>
              ) : (
                <Stack gap="xs">
                  {bookings.map((b) => (
                    <BookingBlock key={b.bookingId ?? b.bookingReference} scheduleId={train.scheduleId} booking={b} />
                  ))}
                </Stack>
              )}
            </Box>
          </Table.Td>
        </Table.Tr>
      )}
    </>
  );
}

function BookingBlock({ scheduleId, booking }: { scheduleId: string; booking: BookingGroup }) {
  const { toast } = useToast();
  const queryClient = useQueryClient();
  const [open, setOpen] = useState(false);
  const [selected, setSelected] = useState<string[]>([]);

  const loadedCount = booking.items.filter((i) => i.status === 'LOADED').length;
  const selectable = booking.items.filter((i) => i.loadable);
  const allSelected = selectable.length > 0 && selectable.every((i) => selected.includes(i.id));
  const toggleItem = (id: string) =>
    setSelected((s) => (s.includes(id) ? s.filter((x) => x !== id) : [...s, id]));
  const toggleAll = () =>
    setSelected((s) =>
      allSelected ? s.filter((id) => !selectable.some((i) => i.id === id)) : selectable.map((i) => i.id),
    );

  const loadMutation = useMutation({
    mutationFn: () => warehouseService.loadItemsOntoTrain(scheduleId, selected),
    onSuccess: (r) => {
      queryClient.invalidateQueries({ queryKey: ['train-loadable-items', scheduleId] });
      queryClient.invalidateQueries({ queryKey: ['loadable-trains'] });
      setSelected([]);
      toast({ title: 'Loaded onto train', description: `Loaded ${r.loadedCount}; skipped ${r.skippedCount}.` });
    },
    onError: (e) =>
      toast({ variant: 'destructive', title: 'Load failed', description: extractErrorMessage(e) }),
  });

  return (
    <Paper withBorder radius="sm" p="xs">
      <Group justify="space-between" style={{ cursor: 'pointer' }} onClick={() => setOpen((o) => !o)} wrap="nowrap">
        <Group gap="xs" wrap="nowrap">
          {open ? <ChevronDown size={15} /> : <ChevronRight size={15} />}
          <Text fw={600}>{booking.bookingReference ?? booking.bookingId?.slice(0, 8) ?? '—'}</Text>
          <Text size="sm" c="dimmed">
            {booking.customerName ?? '—'}
          </Text>
        </Group>
        <Group gap="xs" wrap="nowrap">
          <Badge variant="light" color="blue">
            {booking.items.length} item(s)
          </Badge>
          {loadedCount > 0 && (
            <Badge variant="light" color="green">
              {loadedCount} loaded
            </Badge>
          )}
        </Group>
      </Group>

      {open && (
       <>
        <Table striped highlightOnHover verticalSpacing="xs" mt="xs">
          <Table.Thead>
            <Table.Tr>
              <Table.Th w={36}>
                <Checkbox
                  checked={allSelected}
                  indeterminate={!allSelected && selected.length > 0}
                  onChange={toggleAll}
                  disabled={selectable.length === 0}
                />
              </Table.Th>
              <Table.Th>Container / Cargo</Table.Th>
              <Table.Th>Goods</Table.Th>
              <Table.Th>Weight</Table.Th>
              <Table.Th>Stage</Table.Th>
              <Table.Th>Wagon</Table.Th>
              <Table.Th>Inspection</Table.Th>
            </Table.Tr>
          </Table.Thead>
          <Table.Tbody>
            {booking.items.map((i) => (
              <Table.Tr key={i.id}>
                <Table.Td>
                  <Checkbox checked={selected.includes(i.id)} onChange={() => toggleItem(i.id)} disabled={!i.loadable} />
                </Table.Td>
                <Table.Td>
                  <Text fw={600}>{i.containerNumber ?? i.cargoType ?? '—'}</Text>
                </Table.Td>
                <Table.Td>{i.cargoType ?? '—'}</Table.Td>
                <Table.Td>{weight(i.weight)}</Table.Td>
                <Table.Td>
                  <Badge color={STAGE_COLOR[i.status] ?? 'gray'} variant="light">
                    {i.status.replace(/_/g, ' ')}
                  </Badge>
                </Table.Td>
                <Table.Td>
                  {i.wagonNumber ? (
                    <Badge variant="outline" color="indigo">
                      {i.wagonNumber}
                    </Badge>
                  ) : (
                    <Text size="xs" c="red">
                      Not allocated
                    </Text>
                  )}
                </Table.Td>
                <Table.Td>
                  {i.inspectionStatus ? (
                    <Badge size="xs" variant="light" color={i.inspectionStatus === 'PASSED' ? 'green' : 'orange'}>
                      {i.inspectionStatus}
                    </Badge>
                  ) : (
                    '—'
                  )}
                </Table.Td>
              </Table.Tr>
            ))}
          </Table.Tbody>
        </Table>
        <Group justify="space-between" align="center" mt="xs">
          <Text size="xs" c="dimmed">
            {selected.length} selected · only READY_FOR_LOADING items with a wagon can be loaded
          </Text>
          <Button
            size="compact-sm"
            color="edr-green"
            leftSection={<TrainFront size={14} />}
            disabled={selected.length === 0}
            loading={loadMutation.isPending}
            onClick={() => loadMutation.mutate()}
          >
            Load {selected.length || ''} onto train
          </Button>
        </Group>
       </>
      )}
    </Paper>
  );
}