import { useMemo } from "react";
import { useFieldArray, Controller, useForm } from "react-hook-form";
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import { zodResolver } from "@hookform/resolvers/zod";
import { useNavigate, useParams } from "react-router-dom";
import {
  AlertCircle,
  Check,
  LoaderCircle,
  Loader2,
  Package,
  Weight,
  Plus,
  Trash2,
  MapPin,
  Flame,
  Snowflake,
  Truck,
  FileText,
} from "lucide-react";
import {
  Button,
  Field,
  FieldLabel,
  FieldError,
  Input,
  Badge,
  Switch,
  Textarea,
  Separator,
  Skeleton,
} from "@edr/ui-common";
import type { Freight } from "@edr/types";
import { api } from "@/services/api";
import type { CreateBookingPayload } from "@/services/bookings.service";
import {
  BookingFormInputValues,
  bookingFormSchema,
  getRouteDirection,
  initialBookingFormValues,
  type BookingFormValues,
  type RouteDirection,
} from "./new-booking-form/schema";
import {
  SelectField,
  SelectItem,
  AlertBox,
} from "./new-booking-form/shared";

function yardNameFromBooking(yard: { label?: string; code?: string; name?: string } | undefined | null): string {
  return yard?.label ?? yard?.name ?? yard?.code ?? "";
}

function mapBookingToFormValues(
  booking: Freight.IBooking,
  referenceData: Freight.BookingReferenceData,
): BookingFormInputValues {
  const vals: BookingFormInputValues = {
    ...initialBookingFormValues,
    contractType: (booking.contractType?.toLowerCase() as "new" | "renewal") ?? "new",
    previousContractRef: booking.previousContractId ?? "",
    serviceType:
      booking.serviceType === "RAIL_AND_FORWARDING" ? "rail_forwarding" : "rail",
    firstMile: {
      enabled: booking.firstMileEnabled ?? false,
      pickUpAddress: booking.firstMilePickupAddress ?? "",
    },
    lastMile: {
      enabled: booking.lastMileEnabled ?? false,
      deliveryAddress: booking.lastMileDeliveryAddress ?? "",
    },
    equipmentReturn:
      booking.equipmentReturn === "WITH_RETURN" ? "with_return" : "without_return",
    originYard: yardNameFromBooking(booking.originYard),
    destinationYard: yardNameFromBooking(booking.destinationYard),
    cargoType: booking.freightType === "BULK" ? "bulk" : "container",
    cargoWeight: String(booking.cargoTotalWeightVgm ?? ""),
    isHazardous: booking.isHazardous ?? false,
    isRefrigerated: booking.isRefrigerated ?? false,
    shippingLine: (booking as any).shippingLine?.name ?? "",
    consolidationEnabled: booking.allowConsolidation ?? false,
    notes: "",
    termsAccepted: false,
    freightType: "",
    bulkCommoditytype: "",
    containers: [],
  };

  const bookingCargoTypeId = (booking as any).cargoTypeId as string | undefined;
  if (booking.freightType === "BULK" && bookingCargoTypeId) {
    for (const group of referenceData.cargo_type) {
      const child = group.children?.find((c) => c.id === bookingCargoTypeId);
      if (child) {
        vals.freightType = group.code.toLowerCase();
        vals.bulkCommoditytype = child.name;
        break;
      }
    }
  }

  if (booking.freightType === "CONTAINER" && booking.containers && booking.containers.length > 0) {
    vals.containers = booking.containers.map((c) => ({
      type: c.type === "40ft" ? "40ft" : "20ft" as const,
      containerType: "",
      qty: String(c.qty),
      vgm: String(c.vgm),
    }));
  }

  return vals;
}

export default function EditBookingPage() {
  const { id } = useParams<{ id: string }>();
  const navigate = useNavigate();
  const queryClient = useQueryClient();

  const bookingQuery = useQuery(
    api.bookings.get.queryOptions({
      input: { id: id! },
      enabled: !!id,
    }),
  );

  const { data: referenceData } = useQuery(
    api.bookings.referenceData.queryOptions({
      enabled: !!bookingQuery.data,
    }),
  );

  const updateMutation = useMutation({
    mutationFn: (payload: Partial<CreateBookingPayload>) =>
      api.bookings.update.call({ id: id!, dto: payload }),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: api.bookings.list.queryKey() });
      queryClient.invalidateQueries({ queryKey: api.bookings.get.queryKey({ id: id! }) });
      navigate(`/bookings/${id}`);
    },
  });

  const booking = bookingQuery.data;

  const formValues = useMemo((): BookingFormInputValues | undefined => {
    if (!booking || !referenceData) return undefined;
    return mapBookingToFormValues(booking, referenceData);
  }, [booking, referenceData]);

  const form = useForm<BookingFormInputValues, any, BookingFormValues>({
    defaultValues: initialBookingFormValues,
    values: formValues,
    resolver: zodResolver(bookingFormSchema),
    mode: "onChange",
  });

  const originYard = form.watch("originYard");
  const destinationYard = form.watch("destinationYard");
  const serviceType = form.watch("serviceType");
  const firstMileEnabled = form.watch("firstMile.enabled");
  const lastMileEnabled = form.watch("lastMile.enabled");
  const cargoType = form.watch("cargoType");
  const freightType = form.watch("freightType");
  const containers = form.watch("containers");

  const direction: RouteDirection = useMemo(
    () => getRouteDirection(originYard, destinationYard),
    [originYard, destinationYard],
  );

  const { fields, append, remove } = useFieldArray({
    control: form.control,
    name: "containers",
  });

  const yardOptions = useMemo(() => {
    if (!referenceData?.yard) return [];
    return referenceData.yard.map((y) => ({
      value: y.name,
      label: y.name,
      country: y.country,
    }));
  }, [referenceData]);

  const shippingLineOptions = useMemo(() => {
    if (!referenceData?.shipping_line) return [];
    return referenceData.shipping_line.map((sl) => ({
      value: sl.name,
      label: sl.name,
    }));
  }, [referenceData]);

  const freightTypeGroups = useMemo(() => {
    if (!referenceData?.cargo_type) return [];
    return referenceData.cargo_type.filter(
      (g) => g.code !== "CONTAINER",
    );
  }, [referenceData]);

  const commodityOptions = useMemo(() => {
    if (!referenceData?.cargo_type || !freightType) return [];
    const group = referenceData.cargo_type.find(
      (g) => g.code.toLowerCase() === freightType,
    );
    return group?.children?.map((c) => c.name) ?? [];
  }, [referenceData, freightType]);

  const containerTypeOptions = useMemo(() => {
    if (!referenceData?.containers) return [];
    return referenceData.containers.flatMap((group) =>
      group.types.map((t) => t.name),
    );
  }, [referenceData]);

  const directionStyle: Record<string, string> = {
    export: "bg-sky-50 text-sky-800 border-sky-200",
    import: "bg-amber-50 text-amber-800 border-amber-200",
    domestic: "bg-muted text-muted-foreground border-border",
  };
  const directionLabel: Record<string, string> = {
    export: "Export workflow (inside country to outside country)",
    import: "Import workflow (outside country to inside country)",
    domestic: "Domestic corridor",
  };

  const handleSubmit = form.handleSubmit((data) => {
    const yards = referenceData?.yard ?? [];
    const services = referenceData?.service ?? [];
    const shippingLines = referenceData?.shipping_line ?? [];
    const cargoTree = referenceData?.cargo_type ?? [];
    const containerGroups = referenceData?.containers ?? [];

    const findYardId = (name: string): string =>
      yards.find((y) => y.name === name)?.id ?? "";

    const findServiceTypeId = (): string => {
      const code = data.serviceType === "rail" ? "RAIL" : "RAIL_AND_FORWARDING";
      return services.find((s) => s.code === code)?.id ?? services[0]?.id ?? "";
    };

    const findShippingLineId = (name: string): string | undefined =>
      shippingLines.find((l) => l.name === name)?.id;

    const selectedChild =
      data.cargoType !== "container" && data.bulkCommoditytype
        ? cargoTree
          .find((g) => g.code.toLowerCase() === data.freightType)
          ?.children?.find((c) => c.name === data.bulkCommoditytype)
        : undefined;

    const cargoTypeId =
      data.cargoType === "container"
        ? undefined
        : selectedChild?.id ?? "";

    const findContainerTypeId = (name: string): string => {
      for (const group of containerGroups) {
        const ct = group.types.find((t) => t.name === name);
        if (ct) return ct.id;
      }
      return "";
    };

    const totalWeight =
      data.cargoType === "container"
        ? data.containers.reduce(
          (acc, c) => acc + Number(c.vgm || 0) * Number(c.qty || 0),
          0,
        )
        : Number(data.cargoWeight || 0);

    const apiPayload: Partial<CreateBookingPayload> = {
      scheduledDate: new Date().toISOString().slice(0, 10),
      contractType:
        data.contractType.toUpperCase() as CreateBookingPayload["contractType"],
      serviceTypeId: findServiceTypeId(),
      equipmentReturn:
        data.equipmentReturn === "with_return"
          ? "WITH_RETURN"
          : "WITHOUT_RETURN",
      originYardId: findYardId(data.originYard),
      destinationYardId: findYardId(data.destinationYard),
      tradeDirection:
        direction === "export"
          ? "EXPORT"
          : direction === "domestic"
            ? "DOMESTIC"
            : "IMPORT",
      cargoTypeId: data.cargoType === "container" ? undefined : cargoTypeId,
      cargoTotalWeightVgm: totalWeight,
      isHazardous: data.isHazardous,
      paymentCurrency: "USD",
      allowConsolidation: data.consolidationEnabled,
      // @ts-ignore
      freightType:
        data.cargoType === "container"
          ? ("CONTAINER" as const)
          : ("BULK" as const),
      containers:
        data.cargoType === "container"
          ? data.containers.map((c) => ({
            containerTypeId: findContainerTypeId(c.containerType),
            quantity: Number(c.qty || 1),
            vgmPerUnitTons: Number(c.vgm || 0),
          }))
          : [],
      ...(data.previousContractRef
        ? { previousContractId: data.previousContractRef }
        : {}),
      ...(data.contractType === "renewal" && data.previousContractRef
        ? { pnrCode: data.previousContractRef }
        : {}),
      ...(data.serviceType === "rail_forwarding" && data.firstMile.enabled
        ? { firstMilePickupAddress: data.firstMile.pickUpAddress }
        : {}),
      ...(data.serviceType === "rail_forwarding" && data.lastMile.enabled
        ? { lastMileDeliveryAddress: data.lastMile.deliveryAddress }
        : {}),
      ...(data.shippingLine
        ? { shippingLineId: findShippingLineId(data.shippingLine) }
        : {}),
    };

    updateMutation.mutate(apiPayload);
  });

  if (bookingQuery.isLoading) {
    return (
      <div className="flex h-full items-center justify-center">
        <Loader2 className="size-8 animate-spin text-primary" />
      </div>
    );
  }

  if (bookingQuery.isError || !booking) {
    return (
      <div className="container mx-auto p-6">
        <div className="flex flex-col items-center gap-4 rounded-xl border border-destructive/20 bg-destructive/10 p-12 text-center">
          <AlertCircle className="size-8 text-destructive" />
          <h2 className="text-lg font-bold text-foreground">Failed to load booking</h2>
          <Button variant="outline" onClick={() => navigate("/bookings")}>
            Back to Bookings
          </Button>
        </div>
      </div>
    );
  }

  if (!formValues) {
    return (
      <div className="flex h-full items-center justify-center">
        <Loader2 className="size-8 animate-spin text-primary" />
      </div>
    );
  }

  return (
    <form
      id="edit-booking-form"
      className="mx-auto max-w-4xl"
      onSubmit={handleSubmit}
    >
      <div className="px-6 py-8">
        <h1 className="text-2xl font-bold tracking-tight">
          Edit Booking {booking.reference ?? ""}
        </h1>
        <p className="mt-1 text-sm text-muted-foreground">
          Update the booking details below. All changes are saved together.
        </p>

        {updateMutation.isError && (
          <div className="mt-6 flex items-start gap-3 rounded-xl border border-red-200 bg-red-50 p-4 text-sm text-red-800">
            <AlertCircle className="mt-0.5 h-4 w-4 shrink-0 text-red-500" />
            <div>
              <p className="font-semibold">Failed to save changes</p>
              <p className="mt-1 text-red-600">
                {updateMutation.error instanceof Error
                  ? updateMutation.error.message
                  : "An unexpected error occurred. Please try again."}
              </p>
            </div>
          </div>
        )}

        <div className="mt-8 space-y-10">
          {/* ── Section 1: Contract ── */}
          <section className="space-y-4">
            <div>
              <h2 className="text-lg font-bold">Contract</h2>
              <p className="text-sm text-muted-foreground">
                New contract or renewal of an existing one.
              </p>
            </div>
            <div className="grid gap-4 sm:grid-cols-2">
              <Controller
                name="contractType"
                control={form.control}
                render={({ field, fieldState }) => (
                  <SelectField
                    field={field}
                    error={fieldState.error}
                    label="Contract Type *"
                    placeholder="Select contract type..."
                  >
                    <SelectItem value="new">New Contract</SelectItem>
                    <SelectItem value="renewal">Contract Renewal</SelectItem>
                  </SelectField>
                )}
              />


            </div>
          </section>

          <Separator />

          {/* ── Section 2: Service ── */}
          <section className="space-y-4">
            <div>
              <h2 className="text-lg font-bold">Service</h2>
              <p className="text-sm text-muted-foreground">
                Select the service combination and configure trucking options.
              </p>
            </div>

            <div className="grid gap-4 sm:grid-cols-2">
              <Controller
                name="serviceType"
                control={form.control}
                render={({ field, fieldState }) => (
                  <SelectField
                    field={field}
                    error={fieldState.error}
                    label="Service Type *"
                    placeholder="Select service type..."
                  >
                    <SelectItem value="rail">Rail Transport Only</SelectItem>
                    <SelectItem value="rail_forwarding">Logistics (Rail + Forwarding)</SelectItem>
                  </SelectField>
                )}
              />

              <Controller
                name="equipmentReturn"
                control={form.control}
                render={({ field, fieldState }) => (
                  <SelectField
                    field={field}
                    error={fieldState.error}
                    label="Equipment Return"
                    placeholder="Select..."
                  >
                    <SelectItem value="with_return">With Return</SelectItem>
                    <SelectItem value="without_return">Without Return</SelectItem>
                  </SelectField>
                )}
              />
            </div>

            {serviceType === "rail_forwarding" && (
              <div className="divide-y divide-border rounded-xl border border-border">
                <div className="p-4">
                  <Controller
                    name="firstMile.enabled"
                    control={form.control}
                    render={({ field }) => (
                      <div className="flex items-start justify-between gap-4">
                        <div className="flex items-start gap-3">
                          <Truck className="mt-0.5 h-4 w-4 shrink-0 text-muted-foreground" />
                          <div>
                            <p className="text-sm font-medium">First Mile - Pick-up</p>
                            <p className="mt-0.5 text-xs text-muted-foreground">
                              Truck pick-up from your premises to the origin rail yard.
                            </p>
                          </div>
                        </div>
                        <Switch
                          checked={field.value}
                          onCheckedChange={(value) => {
                            field.onChange(value);
                            if (!value) {
                              form.setValue("firstMile.pickUpAddress", "", {
                                shouldDirty: true,
                                shouldValidate: true,
                              });
                            }
                          }}
                        />
                      </div>
                    )}
                  />
                  {firstMileEnabled && (
                    <Controller
                      name="firstMile.pickUpAddress"
                      control={form.control}
                      render={({ field, fieldState }) => (
                        <Field className="mt-3" data-invalid={fieldState.invalid}>
                          <Input
                            {...field}
                            aria-invalid={fieldState.invalid}
                            placeholder="Pick-up address *"
                          />
                          <FieldError errors={[fieldState.error]} />
                        </Field>
                      )}
                    />
                  )}
                </div>

                <div className="p-4">
                  <Controller
                    name="lastMile.enabled"
                    control={form.control}
                    render={({ field }) => (
                      <div className="flex items-start justify-between gap-4">
                        <div className="flex items-start gap-3">
                          <Truck className="mt-0.5 h-4 w-4 shrink-0 text-muted-foreground" />
                          <div>
                            <p className="text-sm font-medium">Last Mile - Delivery</p>
                            <p className="mt-0.5 text-xs text-muted-foreground">
                              Truck delivery from the destination rail yard to the final address.
                            </p>
                          </div>
                        </div>
                        <Switch
                          checked={field.value}
                          onCheckedChange={(value) => {
                            field.onChange(value);
                            if (!value) {
                              form.setValue("lastMile.deliveryAddress", "", {
                                shouldDirty: true,
                                shouldValidate: true,
                              });
                            }
                          }}
                        />
                      </div>
                    )}
                  />
                  {lastMileEnabled && (
                    <Controller
                      name="lastMile.deliveryAddress"
                      control={form.control}
                      render={({ field, fieldState }) => (
                        <Field className="mt-3" data-invalid={fieldState.invalid}>
                          <Input
                            {...field}
                            aria-invalid={fieldState.invalid}
                            placeholder="Delivery address *"
                          />
                          <FieldError errors={[fieldState.error]} />
                        </Field>
                      )}
                    />
                  )}
                </div>

                <div className="p-4">
                  <Controller
                    name="customsClearingEnabled"
                    control={form.control}
                    render={({ field }) => (
                      <div className="flex items-start justify-between gap-4">
                        <div className="flex items-start gap-3">
                          <FileText className="mt-0.5 h-4 w-4 shrink-0 text-muted-foreground" />
                          <div>
                            <p className="text-sm font-medium">Customs Clearing Service</p>
                            <p className="mt-0.5 text-xs text-muted-foreground">
                              EDR handles customs documentation and clearance on your behalf.
                            </p>
                          </div>
                        </div>
                        <Switch
                          checked={field.value}
                          onCheckedChange={field.onChange}
                        />
                      </div>
                    )}
                  />
                </div>
              </div>
            )}
          </section>

          <Separator />

          {/* ── Section 3: Route ── */}
          <section className="space-y-4">
            <div>
              <h2 className="text-lg font-bold">Route</h2>
              <p className="text-sm text-muted-foreground">
                Select the origin and destination yards.
              </p>
            </div>

            <div className="grid gap-4 sm:grid-cols-2">
              <Controller
                name="originYard"
                control={form.control}
                render={({ field, fieldState }) => (
                  <SelectField
                    field={field}
                    error={fieldState.error}
                    label="Origin Yard *"
                    placeholder="Select origin..."
                    disabled={yardOptions.length === 0}
                  >
                    {yardOptions.length === 0 ? (
                      <SelectItem value="__empty" disabled>No yards available</SelectItem>
                    ) : (
                      yardOptions
                        .filter((y) => y.value !== destinationYard)
                        .map((y) => (
                          <SelectItem key={y.value} value={y.value}>
                            {y.label}
                          </SelectItem>
                        ))
                    )}
                  </SelectField>
                )}
              />

              <Controller
                name="destinationYard"
                control={form.control}
                render={({ field, fieldState }) => (
                  <SelectField
                    field={field}
                    error={fieldState.error}
                    label="Destination Yard *"
                    placeholder="Select destination..."
                    disabled={yardOptions.length === 0}
                  >
                    {yardOptions.length === 0 ? (
                      <SelectItem value="__empty" disabled>No yards available</SelectItem>
                    ) : (
                      yardOptions
                        .filter((y) => y.value !== originYard)
                        .map((y) => (
                          <SelectItem key={y.value} value={y.value}>
                            {y.label}
                          </SelectItem>
                        ))
                    )}
                  </SelectField>
                )}
              />
            </div>

            {direction && (
              <div
                className={`flex items-center gap-2 rounded-lg border px-3 py-2 text-xs font-medium ${directionStyle[direction]}`}
              >
                <MapPin className="h-3.5 w-3.5 shrink-0" />
                {directionLabel[direction]}
              </div>
            )}

            {direction && direction !== "domestic" && (
              <Controller
                name="shippingLine"
                control={form.control}
                render={({ field, fieldState }) => (
                  <SelectField
                    field={field}
                    error={fieldState.error}
                    label="Shipping Line"
                    placeholder="Select shipping line..."
                  >
                    {shippingLineOptions.map((sl) => (
                      <SelectItem key={sl.value} value={sl.value}>
                        {sl.label}
                      </SelectItem>
                    ))}
                  </SelectField>
                )}
              />
            )}

            <div className="space-y-0 divide-y divide-border rounded-xl border border-border">
              <div className="flex items-center justify-between px-4 py-3">
                <div className="flex items-center gap-3">
                  <Flame className="h-4 w-4 shrink-0 text-red-500" />
                  <div>
                    <p className="text-sm font-medium">Hazardous Material</p>
                    <p className="text-xs text-muted-foreground">
                      Applies a Hazard Surcharge to the final bill.
                    </p>
                  </div>
                </div>
                <Controller
                  name="isHazardous"
                  control={form.control}
                  render={({ field }) => (
                    <Switch checked={field.value} onCheckedChange={field.onChange} />
                  )}
                />
              </div>
              <div className="flex items-center justify-between px-4 py-3">
                <div className="flex items-center gap-3">
                  <Snowflake className="h-4 w-4 shrink-0 text-sky-500" />
                  <div>
                    <p className="text-sm font-medium">Refrigerated Cargo</p>
                    <p className="text-xs text-muted-foreground">
                      Temperature-controlled transport applies a Refrigerator Surcharge.
                    </p>
                  </div>
                </div>
                <Controller
                  name="isRefrigerated"
                  control={form.control}
                  render={({ field }) => (
                    <Switch checked={field.value} onCheckedChange={field.onChange} />
                  )}
                />
              </div>
            </div>
          </section>

          <Separator />

          {/* ── Section 4: Cargo ── */}
          <section className="space-y-4">
            <div>
              <h2 className="text-lg font-bold">Cargo Details</h2>
              <p className="text-sm text-muted-foreground">
                Define your cargo type, weight, and container configuration.
              </p>
            </div>

            <div className="grid gap-4 sm:grid-cols-2">
              <Controller
                name="cargoType"
                control={form.control}
                render={({ field, fieldState }) => (
                  <SelectField
                    field={field}
                    error={fieldState.error}
                    label="Cargo Type *"
                    placeholder="Select cargo type..."
                  >
                    <SelectItem value="container">Containerized</SelectItem>
                    <SelectItem value="bulk">General Cargo</SelectItem>
                  </SelectField>
                )}
              />

              <Controller
                name="cargoWeight"
                control={form.control}
                render={({ field, fieldState }) => (
                  <Field data-invalid={fieldState.invalid}>
                    <FieldLabel htmlFor="cargoWeight">
                      Total Cargo Weight (Tons) *
                    </FieldLabel>
                    <div className="relative">
                      <Weight className="absolute left-3 top-1/2 h-4 w-4 -translate-y-1/2 text-muted-foreground" />
                      <Input
                        {...field}
                        id="cargoWeight"
                        type="number"
                        aria-invalid={fieldState.invalid}
                        placeholder="0.00"
                        className="pl-9"
                        min="0"
                        step="0.01"
                      />
                    </div>
                    <FieldError errors={[fieldState.error]} />
                  </Field>
                )}
              />
            </div>

            {cargoType === "bulk" && (
              <>
                <div className="grid gap-4 sm:grid-cols-2">
                  <Controller
                    name="freightType"
                    control={form.control}
                    render={({ field, fieldState }) => (
                      <SelectField
                        field={field}
                        error={fieldState.error}
                        label="Freight Type *"
                        placeholder="Select freight type..."
                      >
                        {freightTypeGroups.map((group) => (
                          <SelectItem key={group.code} value={group.code.toLowerCase()}>
                            {group.name}
                          </SelectItem>
                        ))}
                      </SelectField>
                    )}
                  />

                  {freightType && commodityOptions.length > 0 && (
                    <Controller
                      name="bulkCommoditytype"
                      control={form.control}
                      render={({ field, fieldState }) => (
                        <SelectField
                          field={field}
                          error={fieldState.error}
                          label="Commodity *"
                          placeholder="Select commodity..."
                        >
                          {commodityOptions.map((option) => (
                            <SelectItem key={option} value={option}>
                              {option}
                            </SelectItem>
                          ))}
                        </SelectField>
                      )}
                    />
                  )}
                </div>

                <Controller
                  name="consolidationEnabled"
                  control={form.control}
                  render={({ field }) => (
                    <div className="flex items-center justify-between rounded-xl border border-border px-4 py-3">
                      <div>
                        <p className="text-sm font-medium">Allow Consolidation</p>
                        <p className="text-xs text-muted-foreground">
                          Combine shipments to optimize costs.
                        </p>
                      </div>
                      <Switch checked={field.value} onCheckedChange={field.onChange} />
                    </div>
                  )}
                />
              </>
            )}

            {cargoType === "container" && (
              <div className="space-y-4">
                <div className="flex items-center justify-between">
                  <p className="text-xs font-semibold uppercase tracking-wide text-muted-foreground">
                    Containers
                  </p>
                  <Button
                    type="button"
                    variant="outline"
                    size="sm"
                    onClick={() =>
                      append({
                        type: "20ft",
                        containerType: "",
                        qty: "1",
                        vgm: "",
                      })
                    }
                  >
                    <Plus className="mr-1 h-3.5 w-3.5" />
                    Add Container
                  </Button>
                </div>

                {fields.map((field, index) => {
                  const containerType = containers[index]?.type;
                  const vgm = containers[index]?.vgm ?? 0;
                  const alert = (() => {
                    if (containerType === "20ft" && +vgm > 0) {
                      const limit = direction === "export" ? 25 : 20;
                      if (+vgm > limit) {
                        return `VGM ${vgm}t exceeds the ${limit}t ${direction ?? "standard"} limit for a 20ft container. An Overweight Surcharge will apply.`;
                      }
                    }
                    if (containerType === "40ft" && +vgm > 32.5) {
                      return `VGM ${vgm}t exceeds the 32.5t global limit for a 40ft container. An Overweight Surcharge will apply.`;
                    }
                    return null;
                  })();

                  return (
                    <div
                      key={field.id}
                      className="rounded-xl border border-border p-4 space-y-4"
                    >
                      <div className="flex items-center justify-between">
                        <p className="text-sm font-medium">
                          Container {index + 1}
                        </p>
                        {fields.length > 1 && (
                          <button
                            type="button"
                            onClick={() => remove(index)}
                            className="flex items-center gap-1 text-xs text-destructive hover:underline"
                          >
                            <Trash2 className="h-3.5 w-3.5" />
                            Remove
                          </button>
                        )}
                      </div>

                      <div className="grid gap-3 sm:grid-cols-4">
                        <Controller
                          name={`containers.${index}.type`}
                          control={form.control}
                          render={({ field: typeField, fieldState }) => (
                            <SelectField
                              field={typeField}
                              error={fieldState.error}
                              label="Size *"
                              placeholder="Size..."
                            >
                              <SelectItem value="20ft">20ft (TEU)</SelectItem>
                              <SelectItem value="40ft">40ft (FEU)</SelectItem>
                            </SelectField>
                          )}
                        />

                        <Controller
                          name={`containers.${index}.containerType`}
                          control={form.control}
                          render={({ field: ctField, fieldState }) => (
                            <SelectField
                              field={ctField}
                              error={fieldState.error}
                              label="Type *"
                              placeholder="Type..."
                            >
                              {containerTypeOptions.map((option) => (
                                <SelectItem key={option} value={option}>
                                  {option}
                                </SelectItem>
                              ))}
                            </SelectField>
                          )}
                        />

                        <Controller
                          name={`containers.${index}.qty`}
                          control={form.control}
                          render={({ field: qtyField, fieldState }) => (
                            <Field data-invalid={fieldState.invalid}>
                              <FieldLabel>Quantity *</FieldLabel>
                              <Input
                                value={qtyField.value ?? 1}
                                onChange={(e) => qtyField.onChange(e.target.value)}
                                onBlur={qtyField.onBlur}
                                type="number"
                                aria-invalid={fieldState.invalid}
                                min="1"
                              />
                              <FieldError errors={[fieldState.error]} />
                            </Field>
                          )}
                        />

                        <Controller
                          name={`containers.${index}.vgm`}
                          control={form.control}
                          render={({ field: vgmField, fieldState }) => (
                            <Field data-invalid={fieldState.invalid}>
                              <FieldLabel>VGM (Tons) *</FieldLabel>
                              <Input
                                value={vgmField.value ?? 0}
                                onChange={(e) => vgmField.onChange(e.target.value)}
                                onBlur={vgmField.onBlur}
                                type="number"
                                aria-invalid={fieldState.invalid}
                                placeholder="e.g. 18.5"
                                min="0"
                                step="0.1"
                              />
                              <FieldError errors={[fieldState.error]} />
                            </Field>
                          )}
                        />
                      </div>

                      {alert && (
                        <AlertBox tone="warning">
                          <strong>Overweight Alert:</strong> {alert}
                        </AlertBox>
                      )}
                    </div>
                  );
                })}

                {containers && (() => {
                  const Ft40Wagons = containers
                    .filter((c) => c.type === "40ft")
                    .reduce((sum, c) => sum + Number(c.qty), 0);
                  const Ft20Wagons = containers
                    .filter((c) => c.type === "20ft")
                    .reduce((sum, c) => sum + Number(c.qty), 0);
                  const hasOddUnit = Ft20Wagons % 2 === 1;
                  if (hasOddUnit) {
                    return (
                      <AlertBox tone="warning">
                        <div className="flex items-start gap-2">
                          <div>
                            <p className="font-semibold">Unpaired 20ft Container</p>
                            <p className="mt-1 text-xs">
                              One 20ft container occupies only half a wagon. The wagon
                              will depart once a co-loader is found to fill the
                              remaining slot, which{" "}
                              <strong>may delay departure</strong> beyond the standard
                              lead time.
                            </p>
                          </div>
                        </div>
                      </AlertBox>
                    );
                  }
                  return null;
                })()}
              </div>
            )}
          </section>

          <Separator />

          {/* ── Section 5: Notes & Submit ── */}
          <section className="space-y-4">
            <div>
              <h2 className="text-lg font-bold">Notes & Confirmation</h2>
              <p className="text-sm text-muted-foreground">
                Add any special instructions and confirm the changes.
              </p>
            </div>

            <Controller
              name="notes"
              control={form.control}
              render={({ field }) => (
                <Field>
                  <FieldLabel htmlFor="notes">Additional Notes</FieldLabel>
                  <Textarea
                    {...field}
                    id="notes"
                    placeholder="Any special instructions or notes for EDR operations..."
                    rows={3}
                  />
                </Field>
              )}
            />

            <Controller
              name="termsAccepted"
              control={form.control}
              render={({ field, fieldState }) => (
                <Field data-invalid={fieldState.invalid}>
                  <label className="flex cursor-pointer items-start gap-3">
                    <button
                      type="button"
                      role="checkbox"
                      aria-checked={field.value}
                      aria-invalid={fieldState.invalid}
                      onClick={() => field.onChange(!field.value)}
                      className={`mt-0.5 flex h-5 w-5 shrink-0 items-center justify-center rounded border-2 transition ${field.value ? "border-primary bg-primary" : "border-input"
                      }`}
                    >
                      {field.value && (
                        <Check className="h-3 w-3 text-primary-foreground" />
                      )}
                    </button>
                    <span className="text-sm text-muted-foreground">
                      I confirm the information is accurate and agree to EDR's{" "}
                      <span className="text-primary">
                        freight contract terms and conditions
                      </span>
                      .
                    </span>
                  </label>
                  <FieldError errors={[fieldState.error, form.formState.errors.termsAccepted]} />
                </Field>
              )}
            />
          </section>
        </div>

        {/* ── Submit ── */}
        <div className="mt-10 flex items-center gap-3 border-t border-border pt-6">
          <Button
            type="button"
            variant="outline"
            onClick={() => navigate(`/bookings/${id}`)}
          >
            Cancel
          </Button>
          <Button
            type="submit"
            form="edit-booking-form"
            disabled={updateMutation.isPending}
          >
            {updateMutation.isPending ? (
              <LoaderCircle className="h-4 w-4 animate-spin" />
            ) : (
              <Check />
            )}
            {updateMutation.isPending ? "Saving Changes..." : "Save Changes"}
          </Button>
        </div>
      </div>
    </form>
  );
}
