import { useMemo, useState } from "react";
import { Badge, Box, Group, Paper, Stack, Text, ThemeIcon } from "@mantine/core";
import { MousePointerClick, TrainFront } from "lucide-react";
import type { TrainScheduleDetail } from "@/types/trainScheduling";
import { TrainStatsBar } from "./TrainStatsBar";
import { WagonCard } from "./WagonCard";
import { InteractiveTrainConsist } from "./InteractiveTrainConsist";
import { RemoveBookingModal } from "./RemoveBookingModal";
import { useMutation } from "@tanstack/react-query";
import { api } from "@/services/api";
import { freightBrand } from "@/theme/freight-brand";

type Wagon = NonNullable<TrainScheduleDetail["trainSet"]>["wagons"][number];

interface TrainConsistViewProps {
  scheduleDetail: TrainScheduleDetail;
  scheduleId: string;
  maxWagons: number;
  /** Booking id selected in the side panel — highlights its wagons in the consist. */
  highlightBookingId?: string | null;
}

function LegendDot({ color, label, dashed }: { color: string; label: string; dashed?: boolean }) {
  return (
    <Group gap={5} wrap="nowrap">
      <Box
        style={{
          width: 10,
          height: 10,
          borderRadius: 3,
          background: dashed ? "var(--mantine-color-gray-1)" : `var(--mantine-color-${color}-5)`,
          border: dashed ? "1.5px dashed var(--mantine-color-gray-4)" : "none",
        }}
      />
      <Text size="xs" c="dimmed">
        {label}
      </Text>
    </Group>
  );
}

export const TrainConsistView = ({
  scheduleDetail,
  scheduleId,
  maxWagons,
  highlightBookingId,
}: TrainConsistViewProps) => {
  const [selectedWagonId, setSelectedWagonId] = useState<string | null>(null);
  const [removeModalOpen, setRemoveModalOpen] = useState(false);

  const unassignMutation = useMutation(
    api.trainScheduling.unassignBooking.mutationOptions(),
  );
  const removeWagonMutation = useMutation(
    api.trainScheduling.removeWagonSlot.mutationOptions(),
  );

  const trainSet = scheduleDetail.trainSet;
  const wagons = trainSet?.wagons ?? [];

  // Join company/customer name from schedule bookings by booking id.
  const companyByBooking = useMemo(() => {
    const map = new Map<string, string>();
    for (const b of scheduleDetail.bookings ?? []) {
      if (b.id && b.customer) map.set(b.id, b.customer);
    }
    return map;
  }, [scheduleDetail.bookings]);

  const selectedWagon = wagons.find((w) => w.id === selectedWagonId) ?? null;
  const loadedCount = wagons.filter((w) => (w.allocations?.length ?? 0) > 0).length;

  const handleRemoveBooking = (wagon: Wagon) => {
    setSelectedWagonId(wagon.id);
    setRemoveModalOpen(true);
  };

  const handleConfirmRemoveBooking = async () => {
    if (selectedWagon?.allocations?.[0]?.bookingId) {
      await unassignMutation.mutateAsync({
        id: scheduleId,
        bookingId: selectedWagon.allocations[0].bookingId,
      });
      setRemoveModalOpen(false);
      setSelectedWagonId(null);
    }
  };

  const handleRemoveWagon = async (wagonId: string) => {
    if (confirm("Are you sure you want to remove this wagon slot?")) {
      await removeWagonMutation.mutateAsync({ scheduleId, wagonId });
      setSelectedWagonId(null);
    }
  };

  const weightUsed = wagons.reduce(
    (sum, w) => sum + (w.allocations?.[0]?.allocatedWeightTons ?? 0),
    0,
  );
  const lengthUsed = wagons.reduce((sum, w) => sum + (w.lengthMeters ?? 0), 0);

  return (
    <Stack gap="md" style={{ width: "100%" }}>
      <TrainStatsBar
        weightUsed={weightUsed}
        weightMax={trainSet?.locomotive?.maxPullWeightTons ?? null}
        lengthUsed={lengthUsed}
        lengthMax={trainSet?.locomotive?.maxTrainLengthMeters ?? null}
        wagonCount={wagons.length}
        wagonMax={maxWagons}
      />

      {/* Consist panel */}
      <Paper
        radius="lg"
        withBorder
        style={{ borderColor: "var(--mantine-color-gray-2)", overflow: "hidden" }}
      >
        <Group
          justify="space-between"
          wrap="nowrap"
          px="md"
          py="sm"
          style={{ borderBottom: "1px solid var(--mantine-color-gray-2)" }}
        >
          <Group gap={10} wrap="nowrap">
            <Box
              style={{
                width: 34,
                height: 34,
                borderRadius: 9,
                display: "flex",
                alignItems: "center",
                justifyContent: "center",
                background: freightBrand.gradient,
                color: "white",
              }}
            >
              <TrainFront size={18} />
            </Box>
            <div>
              <Text fw={800} size="sm">
                Train consist
              </Text>
              <Text size="11px" c="dimmed">
                {wagons.length} wagons · {loadedCount} loaded · {wagons.length - loadedCount} empty
              </Text>
            </div>
          </Group>
          <Group gap="md" wrap="nowrap" visibleFrom="sm">
            <LegendDot color="cyan" label="Container" />
            <LegendDot color="orange" label="Bulk" />
            <LegendDot color="gray" label="Empty" dashed />
          </Group>
        </Group>

        <Box p="md">
          <InteractiveTrainConsist
            wagons={wagons}
            locomotive={trainSet?.locomotive}
            getCompany={(bookingId) => (bookingId ? companyByBooking.get(bookingId) ?? null : null)}
            selectedWagonId={selectedWagonId}
            onSelectWagon={(w) => setSelectedWagonId((prev) => (prev === w.id ? null : w.id))}
            highlightBookingId={highlightBookingId}
          />
        </Box>
      </Paper>

      {/* Selected wagon — editable detail card */}
      {selectedWagon ? (
        <Box>
          <Group gap={6} mb={6} wrap="nowrap">
            <Badge variant="light" color="edr-green" radius="sm">
              Editing wagon #{selectedWagon.sequenceNo}
            </Badge>
            <Text size="xs" c="dimmed">
              Update container numbers or remove the booking
            </Text>
          </Group>
          <WagonCard
            wagon={selectedWagon}
            company={
              selectedWagon.allocations?.[0]?.bookingId
                ? companyByBooking.get(selectedWagon.allocations[0].bookingId) ?? null
                : null
            }
            scheduleId={scheduleId}
            scheduleStatus={scheduleDetail.status}
            onRemoveBooking={handleRemoveBooking}
            onRemoveWagon={handleRemoveWagon}
          />
        </Box>
      ) : wagons.length ? (
        <Paper
          radius="md"
          py="sm"
          px="md"
          style={{
            background: "var(--mantine-color-gray-0)",
            border: "1px dashed var(--mantine-color-gray-3)",
          }}
        >
          <Group gap={8} justify="center" c="dimmed">
            <ThemeIcon size={24} radius="xl" variant="light" color="gray">
              <MousePointerClick size={13} />
            </ThemeIcon>
            <Text size="xs" c="dimmed">
              Click a wagon in the train to edit container numbers or remove its booking.
            </Text>
          </Group>
        </Paper>
      ) : null}

      <RemoveBookingModal
        opened={removeModalOpen}
        onClose={() => {
          setRemoveModalOpen(false);
        }}
        wagon={selectedWagon}
        onConfirm={handleConfirmRemoveBooking}
        isLoading={unassignMutation.isPending}
      />
    </Stack>
  );
};
