import { attestationHeaders } from "./attestation.js";
import type { WorkerRecord } from "./keys.js";
import { fetchHeld, readResponseJson } from "./http.js";
import type { Transport } from "./proxy.js";

const STATS_TIMEOUT_MS = 3000;
const STATS_CONCURRENCY = 8;

export type WorkerStatsResult = {
  sessionId: string;
  clientId: string;
  hostname: string;
} & ({ stats: Record<string, unknown> } | { error: "unreachable" });

/** Collects `GET <advertiseUrl>/_pinned/stats` from every worker, attested per
 * hop; a failing or slow worker reports `error: "unreachable"`. */
export const collectWorkerStats = async (
  workers: WorkerRecord[],
  now: number,
  fetchImpl: Transport.FetchLike,
): Promise<WorkerStatsResult[]> => {
  const results: WorkerStatsResult[] = new Array<WorkerStatsResult>(
    workers.length,
  );
  let next = 0;
  const run = async () => {
    while (next < workers.length) {
      const index = next++;
      const worker = workers[index]!;
      const identity = {
        sessionId: worker.sessionId,
        clientId: worker.clientId,
        hostname: worker.hostname,
      };
      try {
        const res = await fetchHeld(
          fetchImpl,
          new Request(`${worker.advertiseUrl}/_pinned/stats`, {
            headers: await attestationHeaders(worker, {}, now),
            redirect: "manual",
            signal: AbortSignal.timeout(STATS_TIMEOUT_MS),
          }),
        );
        if (!res.ok) throw new Error(`HTTP ${res.status}`);
        const stats = await readResponseJson(res);
        if (typeof stats !== "object" || stats === null) {
          throw new Error("malformed stats");
        }
        const scoped = stats as Record<string, unknown>;
        const namespaces = new Set(worker.namespaces);
        const rawNamespaces = scoped["namespaces"];
        const rawInstances = scoped["instances"];
        results[index] = {
          ...identity,
          stats: {
            ...scoped,
            ...(typeof rawNamespaces === "object" &&
            rawNamespaces !== null &&
            !Array.isArray(rawNamespaces)
              ? {
                  namespaces: Object.fromEntries(
                    Object.entries(rawNamespaces).filter(([namespace]) =>
                      namespaces.has(namespace),
                    ),
                  ),
                }
              : {}),
            ...(Array.isArray(rawInstances)
              ? {
                  instances: rawInstances.filter(
                    (instance) =>
                      typeof instance === "object" &&
                      instance !== null &&
                      namespaces.has(
                        (instance as Record<string, unknown>)[
                          "namespace"
                        ] as string,
                      ),
                  ),
                }
              : {}),
          },
        };
      } catch {
        results[index] = { ...identity, error: "unreachable" };
      }
    }
  };
  await Promise.all(
    Array.from({ length: Math.min(STATS_CONCURRENCY, workers.length) }, run),
  );
  return results;
};