import { useQuery, useQueries } from "@tanstack/react-query";
import {
  getUnitList,
  getChildUnits,
} from "@/user-management/services/api/unitService";

const extractItems = (data: any): any[] => {
  if (!data) return [];
  if (Array.isArray(data)) return data;
  if (Array.isArray(data.items)) return data.items;
  return [];
};

/**
 * Total count of every unit in an organization's hierarchy:
 * top-level units from `/units/list/{orgId}` plus every descendant
 * reachable via `/units/child-units/{parentId}`.
 *
 * Both create-unit and relate-from-sub-city flows count against this
 * same total so the configured cap is enforced consistently.
 */
export const useOrgUnitTotalCount = (organizationId: string) => {
  const topLevelQuery = useQuery({
    queryKey: ["unitList", organizationId, "capacity-total"],
    queryFn: () => getUnitList(organizationId, { take: 1000, skip: 0 }),
    enabled: !!organizationId,
    staleTime: 0,
  });

  const topLevelItems = extractItems(topLevelQuery.data?.data);
  const topLevelIds: string[] = topLevelItems
    .map((u: any) => u?.id)
    .filter(Boolean);

  const childQueries = useQueries({
    queries: topLevelIds.map((id) => ({
      queryKey: ["unitChildren", id, "capacity-total"],
      queryFn: () => getChildUnits(id),
      enabled: !!id,
      staleTime: 0,
    })),
  });

  // De-dupe units by `id` because some backends include the parent in
  // `/units/child-units/{parentId}` responses, and some `/units/list/{orgId}`
  // implementations already include non-top-level units.
  const uniqueIds = new Set<string>();
  for (const unit of topLevelItems) {
    const id = unit?.id;
    if (id) uniqueIds.add(String(id));
  }
  for (const q of childQueries) {
    if (!q.data) continue;
    const items = extractItems((q.data as any).data);
    for (const unit of items) {
      const id = unit?.id;
      if (id) uniqueIds.add(String(id));
    }
  }

  const total = uniqueIds.size;
  const isLoading =
    topLevelQuery.isLoading ||
    topLevelQuery.isFetching ||
    childQueries.some((q) => q.isLoading || q.isFetching);

  return { total, isLoading };
};
