emaui/apps/portal/src/app/features/location/components/LocationPicker.tsx

import { useState, useMemo, useEffect, useCallback } from 'react';
import { Stack, Select, Group, Text, Loader, Center, Badge } from '@mantine/core';
import { ErrorState } from '@ema-platform/ui';
import { useLocalized } from '@ema-platform/api';
import { useGetLocationTypesQuery, useGetLocationsQuery } from '../api/location-api';
import type { Location, LocationType } from '../types/location';
import { useTranslation } from 'react-i18next';

interface LocationPickerProps {
  value?: string;
  onChange?: (locationId: string | null) => void;
  onChainChange?: (chain: Location[]) => void;
  required?: boolean;
  /** Caps how many cascading levels render (e.g. 3 = City / Sub-city / Woreda only). */
  maxDepth?: number;
}

export function LocationPicker({ value, onChange, onChainChange, required, maxDepth }: LocationPickerProps) {
  const { t } = useTranslation();
  const localized = useLocalized();
  const { data: typesRes, isLoading: typesLoading, isError: typesError, refetch: refetchTypes } = useGetLocationTypesQuery();
  const { data: locsRes, isLoading: locsLoading, isError: locsError, refetch: refetchLocs } = useGetLocationsQuery({ take: 10000 });

  const locationTypes = typesRes?.items ?? [];
  const allLocations = locsRes?.items ?? [];

  const locMap = useMemo(() => {
    const map = new Map<string, Location>();
    allLocations.forEach((loc) => map.set(loc.id, loc));
    return map;
  }, [allLocations]);

  const typeMap = useMemo(() => {
    const map = new Map<string, LocationType>();
    locationTypes.forEach((t) => map.set(t.id, t));
    return map;
  }, [locationTypes]);

  const childrenByParentId = useMemo(() => {
    const map = new Map<string, Location[]>();
    allLocations.forEach((loc) => {
      const key = loc.parentId ?? '__root__';
      if (!map.has(key)) map.set(key, []);
      map.get(key)!.push(loc);
    });
    return map;
  }, [allLocations]);

  const [selectedChain, setSelectedChain] = useState<Location[]>([]);

  useEffect(() => {
    if (!value || locMap.size === 0) return;
    const loc = locMap.get(value);
    if (!loc) return;

    const chain: Location[] = [];
    let current: Location | undefined = loc;
    while (current) {
      chain.unshift(current);
      current = current.parentId ? locMap.get(current.parentId) : undefined;
    }
    setSelectedChain(chain);
    onChainChange?.(chain);
  }, [value, locMap, onChainChange]);

  const currentLevelChildren = useMemo(() => {
    const parentId =
      selectedChain.length === 0
        ? null
        : selectedChain[selectedChain.length - 1].id;
    const key = parentId ?? '__root__';
    return childrenByParentId.get(key) ?? [];
  }, [selectedChain, childrenByParentId]);

  const levelLabel = useMemo(() => {
    if (currentLevelChildren.length === 0) return '';
    const typeIds = [...new Set(currentLevelChildren.map((c) => c.locationTypeId))];
    const names = typeIds
      .map((id) => localized(typeMap.get(id)?.names))
      .filter(Boolean);
    return names.join(' / ');
  }, [currentLevelChildren, typeMap, localized]);

  const depth = selectedChain.length;

  const allLevelsComplete = useMemo(() => {
    return selectedChain.every((loc, i) => {
      const key = loc.id;
      const children = childrenByParentId.get(key);
      return !children || children.length === 0;
    });
  }, [selectedChain, childrenByParentId]);

  const handleSelect = useCallback(
    (id: string | null) => {
      if (!id) {
        const newChain = selectedChain.slice(0, -1);
        setSelectedChain(newChain);
        onChange?.(newChain.length > 0 ? newChain[newChain.length - 1].id : null);
        onChainChange?.(newChain);
        return;
      }

      const loc = locMap.get(id);
      if (!loc) return;

      const newChain = selectedChain.slice(0, depth);
      newChain.push(loc);
      setSelectedChain(newChain);

      onChange?.(id);
      onChainChange?.(newChain);
    },
    [selectedChain, locMap, onChange, onChainChange, depth],
  );

  const buildOptions = (levelIdx: number) => {
    if (levelIdx === 0) {
      const roots = childrenByParentId.get('__root__') ?? [];
      return roots
        .map((loc) => ({ value: loc.id, label: localized(loc.names) }))
        .sort((a, b) => a.label.localeCompare(b.label));
    }

    const parent = selectedChain[levelIdx - 1];
    if (!parent) return [];
    const children = childrenByParentId.get(parent.id) ?? [];
    return children
      .map((loc) => ({ value: loc.id, label: localized(loc.names) }))
      .sort((a, b) => a.label.localeCompare(b.label));
  };

  const getLevelLabel = (levelIdx: number) => {
    if (levelIdx === 0) {
      const roots = childrenByParentId.get('__root__') ?? [];
      if (roots.length === 0) return '';
      const typeIds = [...new Set(roots.map((r) => r.locationTypeId))];
      const names = typeIds
        .map((id) => localized(typeMap.get(id)?.names))
        .filter(Boolean);
      return names.join(' / ');
    }

    const parent = selectedChain[levelIdx - 1];
    if (!parent) return t('location.chooseFirst');
    const children = childrenByParentId.get(parent.id) ?? [];
    const typeIds = [...new Set(children.map((c) => c.locationTypeId))];
    const names = typeIds
      .map((id) => localized(typeMap.get(id)?.names))
      .filter(Boolean);
    return names.join(' / ');
  };

  const selectedPath = useMemo(() => {
    return selectedChain
      .map((loc) => localized(loc.names))
      .join(' → ');
  }, [selectedChain, localized]);

  if (typesLoading || locsLoading) {
    return (
      <Center py="md">
        <Loader size="sm" />
      </Center>
    );
  }

  if (typesError || locsError) {
    return (
      <ErrorState
        title={t('location.loadFailed')}
        onRetry={() => {
          refetchTypes();
          refetchLocs();
        }}
      />
    );
  }

  const roots = childrenByParentId.get('__root__') ?? [];

  if (roots.length === 0) {
    return (
      <Text size="sm" c="dimmed">
        {t('location.noLocationsAvailable')}
      </Text>
    );
  }

  // Only offer a further level when its location type resolves to a known
  // name — an unnamed/unmapped type (e.g. a stray Kebele row) would otherwise
  // render as a dead-end "Sub-location" picker whose selection is silently
  // dropped (AddressFormContent only maps Region/City/SubCity/Woreda).
  const totalRenderedLevels = Math.min(
    maxDepth ?? Infinity,
    Math.max(1, selectedChain.length + (currentLevelChildren.length > 0 && levelLabel ? 1 : 0)),
  );

  const levels = Array.from({ length: totalRenderedLevels }, (_, i) => i);

  return (
    <Stack gap="sm">
      <Group gap="xs" align="end" wrap="wrap">
        {levels.map((levelIdx) => {
          const options = buildOptions(levelIdx);
          const currentValue = selectedChain[levelIdx]?.id ?? null;
          const isDisabled = levelIdx > 0 && !selectedChain[levelIdx - 1];

          return (
            <Select
              key={levelIdx}
              label={getLevelLabel(levelIdx) || `Level ${levelIdx + 1}`}
              placeholder={t('location.select')}
              data={options}
              value={currentValue}
              onChange={(val) => handleSelect(val)}
              disabled={isDisabled}
              searchable
              clearable
              size="sm"
              style={{ minWidth: 160, flex: 1 }}
              nothingFoundMessage={t('location.noOptions')}
              required={required && levelIdx === levels.length - 1}
            />
          );
        })}
      </Group>

      {selectedPath && (
        <Group gap="xs">
          {selectedChain.map((loc) => {
            const typeInfo = typeMap.get(loc.locationTypeId);
            return (
              <Badge
                key={loc.id}
                size="sm"
                variant="light"
                color="blue"
                style={{ textTransform: 'none' }}
              >
                {typeInfo ? `${localized(typeInfo.names)}: ` : ''}
                {localized(loc.names)}
              </Badge>
            );
          })}
        </Group>
      )}
    </Stack>
  );
}