import { useState, useMemo } from "react";
import { Plus } from "lucide-react";
import { Button, Group, Modal, NumberInput, Select, Stack, Text } from "@mantine/core";

import { Freight } from "@edr/types";

import { useMutation, useQuery } from "@tanstack/react-query";

import { api } from "@/services/api";
import { useToast } from "@/hooks/use-toast";

export function AssignWagonDialog({ trainId }: { trainId: string }) {
  const [open, setOpen] = useState(false);
  const [wagonId, setWagonId] = useState<string | null>(null);
  const [sequence, setSequence] = useState<number | "">("");
  const { data: wagons } = useQuery(api.wagons.list.queryOptions({ input: {} }));
  const { data: yards = [] } = useQuery(api.routes.yards.queryOptions());
  const assign = useMutation(api.wagons.assignToTrain.mutationOptions());
  const { toast } = useToast();

  const available = (wagons ?? []).filter(
    (w) => w.status === Freight.WagonStatus.Available || !w.trainId,
  );

  const yardLabelById = useMemo(
    () =>
      new Map(
        yards.map((y) => [y.id, y.label ?? y.code ?? y.id]),
      ),
    [yards],
  );

  const wagonOptions = available.map((w) => ({
    value: w.id,
    label: w.currentYardId
      ? `${w.wagonNumber} (${yardLabelById.get(w.currentYardId) ?? "yard"})`
      : `${w.wagonNumber} (no yard)`,
  }));

  const handleAssign = async () => {
    if (!wagonId) return;
    try {
      await assign.mutateAsync({
        wagonId,
        trainId,
        sequenceNumber: sequence === "" ? undefined : Number(sequence),
      });
      toast({ title: "Wagon attached to train" });
      setOpen(false);
      setWagonId(null);
      setSequence("");
    } catch {
      toast({ title: "Failed to assign wagon", variant: "destructive" });
    }
  };

  return (
    <>
      <Button
        color="edr-green"
        size="sm"
        radius="lg"
        leftSection={<Plus size={16} />}
        onClick={() => setOpen(true)}
      >
        Assign wagon
      </Button>

      <Modal
        opened={open}
        onClose={() => setOpen(false)}
        title={<Text fw={600}>Assign wagon to train</Text>}
        radius="lg"
        centered
      >
        <Stack gap="md">
          <Select
            label="Wagon"
            placeholder="Select wagon"
            data={wagonOptions}
            value={wagonId}
            onChange={setWagonId}
            searchable
          />
          <NumberInput
            label="Sequence (optional)"
            value={sequence}
            onChange={(value) => setSequence(value === "" ? "" : Number(value))}
            min={1}
          />
          <Group justify="flex-end">
            <Button variant="default" onClick={() => setOpen(false)}>
              Cancel
            </Button>
            <Button color="edr-green" loading={assign.isPending} onClick={handleAssign}>
              Assign
            </Button>
          </Group>
        </Stack>
      </Modal>
    </>
  );
}