/** Backend-agnostic diagnose contract.
 *
 * Two providers, one shape:
 * - "zerogpu": Gradio Space endpoint /diagnose via @gradio/client.
 * - "vastai":  POST ${url}/score multipart (ct_files) -> {scores, rows}.
 * No weights or API keys here; the Vast key never leaves the server.
 */

import { Client } from "@gradio/client";

export type BackendKind = "zerogpu" | "vastai";

export interface FindingRow {
  finding: string;
  score: number;
}

export interface DiagnoseResult {
  scores: Record<string, number>;
  rows: FindingRow[];
}

export function backendKind(): BackendKind {
  const raw = (import.meta.env.PUBLIC_BACKEND_KIND ?? "zerogpu").toLowerCase();
  if (raw !== "zerogpu" && raw !== "vastai") throw new Error("Unknown backend");
  return raw;
}

export function backendLabel(): string {
  return backendKind() === "vastai" ? "Vast.ai" : "ZeroGPU";
}

/** Gradio Dataframe payload varies by version: accept {data}, {headers,data}, or raw rows. */
export function parseDataframe(payload: unknown): FindingRow[] {
  let data: unknown = payload;
  if (payload !== null && typeof payload === "object" && "data" in payload) {
    data = (payload as { data: unknown }).data;
  }
  if (!Array.isArray(data)) return [];
  const rows: FindingRow[] = [];
  for (const entry of data) {
    if (Array.isArray(entry) && entry.length >= 2) {
      // Null/empty scores must be skipped BEFORE Number(): Number(null) is
      // 0, which would fabricate a zero-scored finding (NaN cells come back
      // null over the API for empty CSV cells).
      if (entry[1] === null || entry[1] === undefined || entry[1] === "") continue;
      const score = Number(entry[1]);
      if (entry[0] != null && entry[0] !== "" && Number.isFinite(score)) {
        rows.push({ finding: String(entry[0]), score });
      }
    }
  }
  return rows;
}

/** gr.Label arrives as {label, confidences}; older clients may send a flat dict. */
export function parseLabel(payload: unknown): Record<string, number> {
  if (payload === null || typeof payload !== "object") return {};
  const confidences = (payload as { confidences?: unknown }).confidences;
  if (Array.isArray(confidences)) {
    const out: Record<string, number> = {};
    for (const entry of confidences) {
      const label = (entry as { label?: unknown }).label;
      const raw = (entry as { confidence?: unknown }).confidence;
      if (raw === null || raw === undefined || raw === "") continue;
      const confidence = Number(raw);
      if (typeof label === "string" && label !== "" && Number.isFinite(confidence)) {
        out[label] = confidence;
      }
    }
    return out;
  }
  const rec = payload as Record<string, unknown>;
  const vals = Object.values(rec);
  if (
    vals.length > 0 &&
    vals.every((v) => v !== null && v !== "" && Number.isFinite(Number(v)))
  ) {
    const out: Record<string, number> = {};
    for (const [k, v] of Object.entries(rec)) out[k] = Number(v);
    return out;
  }
  return {};
}

export function normalize(
  scores: Record<string, number>,
  rows: FindingRow[],
): DiagnoseResult {
  const clean: Record<string, number> = {};
  for (const [k, raw] of Object.entries(scores ?? {})) {
    if (raw === null || raw === undefined) continue;
    const v: unknown = raw;
    if (typeof v === "string" && v.trim() === "") continue;
    const n = Number(v);
    if (Number.isFinite(n)) clean[k] = n;
  }
  const sorted = [...rows].sort((a, b) => b.score - a.score);
  return { scores: clean, rows: sorted };
}

export async function diagnoseRemote(files: File[]): Promise<DiagnoseResult> {
  const kind = backendKind();
  const url = import.meta.env.PUBLIC_BACKEND_URL ?? "scomb2/Radar";

  if (kind === "vastai") {
    const form = new FormData();
    for (const f of files) form.append("ct_files", f);
    let res: Response;
    try {
      res = await fetch(`${url}/score`, { method: "POST", body: form });
    } catch {
      throw new Error("GPU busy or unreachable — Retry");
    }
    if (!res.ok) {
      const body = (await res.json().catch(() => null)) as { detail?: string } | null;
      throw new Error(body?.detail ?? "GPU busy or unreachable — Retry");
    }
    const json = (await res.json()) as { scores: Record<string, number>; rows: FindingRow[] };
    return normalize(json.scores, json.rows ?? []);
  }

  // zerogpu: stable /diagnose endpoint on the headless Space.
  let client: Client;
  try {
    client = await Client.connect(url);
  } catch {
    throw new Error("GPU busy or unreachable — Retry");
  }
  let data: unknown[];
  try {
    // Upload the original Files first: the client's payload walker rewraps
    // every File as a nameless Blob, and server intake routes on the file
    // extension — so predict() receives explicit FileData with orig_name.
    // (upload_files/config are untyped in this client version; the lockfile
    // pins the exact build this was verified against.)
    const root = client.config?.root ?? url;
    const uploaded = await client.upload_files(root, files);
    const paths: string[] = uploaded?.files ?? [];
    if (paths.length !== files.length) {
      throw new Error("GPU busy or unreachable — Retry");
    }
    const result = await client.predict("/diagnose", {
      ct_files: paths.map((path, i) => ({
        path,
        orig_name: files[i].name,
        meta: { _type: "gradio.FileData" },
      })),
    });
    data = result.data as unknown[];
  } catch (e) {
    // @gradio/client rejects with a status object, not always an Error —
    // read message structurally so gr.Error text still renders verbatim.
    const message = (e as { message?: unknown } | null | undefined)?.message;
    if (typeof message === "string" && message) throw new Error(message);
    throw new Error("GPU busy or unreachable — Retry");
  }
  const rows = parseDataframe(data[1]);
  return normalize(parseLabel(data[0]), rows);
}