import { removeBackground } from "@imgly/background-removal";
import type { ToolResult } from "../../core/io/output-types";
import type { ToolProgress } from "../../core/plugins/plugin-types";
import { installImglyCache, isImglyAssetCached } from "./imgly-cache";

export interface BackgroundRemoverOptions {
  /** Quality/size trade-off: small=quint8, medium=fp16, large=fp32. */
  quality: "small" | "medium" | "large";
  /** Use WebGPU when available (falls back to wasm automatically). */
  useGpu: boolean;
}

/**
 * Self-hosted asset location. The @imgly model + onnxruntime-web wasm files are
 * mirrored into `public/imgly/` (see scripts/fetch-imgly-assets.mjs) so nothing
 * is fetched from a third party — the image never leaves the browser.
 */
function publicPath(origin: string): string {
  return new URL("/imgly/", origin).toString();
}

const MODEL_BY_QUALITY: Record<
  BackgroundRemoverOptions["quality"],
  "isnet" | "isnet_fp16" | "isnet_quint8"
> = {
  small: "isnet_quint8",
  medium: "isnet_fp16",
  large: "isnet",
};

interface ResourceEntry {
  size: number;
  chunks: Array<{ name: string }>;
}
type ResourceMap = Record<string, ResourceEntry>;

function resourceKeys(options: BackgroundRemoverOptions): string[] {
  const wasmBase = options.useGpu
    ? "/onnxruntime-web/ort-wasm-simd-threaded.jsep"
    : "/onnxruntime-web/ort-wasm-simd-threaded";

  return [
    `/models/${MODEL_BY_QUALITY[options.quality]}`,
    `${wasmBase}.wasm`,
    `${wasmBase}.mjs`,
  ];
}

async function loadResourceMap(base: string): Promise<ResourceMap | null> {
  try {
    const res = await fetch(new URL("resources.json", base).toString());
    if (!res.ok) return null;
    return (await res.json()) as ResourceMap;
  } catch {
    return null;
  }
}

export function createBackgroundRemovalProgress(
  expectedSizes: ReadonlyMap<string, number>,
  reportProgress?: (progress: ToolProgress) => void
): (key: string, current: number, total: number) => void {
  const fetchedByKey = new Map<string, number>();
  const expectedTotal = [...expectedSizes.values()].reduce(
    (sum, size) => sum + size,
    0
  );
  let sawFetch = false;
  let lastPercentage = 0;

  return (key, current, total) => {
    let percentage: number;
    let message: string;

    if (key.startsWith("fetch:")) {
      sawFetch = true;
      const resourceKey = key.slice("fetch:".length);
      fetchedByKey.set(resourceKey, Math.min(current, total));

      const fetched = [...fetchedByKey.values()].reduce(
        (sum, size) => sum + size,
        0
      );
      const knownTotal =
        expectedTotal > 0
          ? expectedTotal
          : [...fetchedByKey.keys()].reduce(
              (sum, knownKey) =>
                sum + (knownKey === resourceKey ? total : fetchedByKey.get(knownKey) ?? 0),
              0
            );

      percentage = knownTotal > 0 ? (fetched / knownTotal) * 65 : lastPercentage;
      message = "Loading AI model...";
    } else if (key.startsWith("compute:")) {
      const computePercentage = total > 0 ? current / total : 0;
      percentage = sawFetch
        ? 65 + computePercentage * 35
        : computePercentage * 100;
      message =
        current === 0 ? "Preparing image..." : "Removing background...";
    } else {
      percentage = lastPercentage;
      message = "Processing...";
    }

    lastPercentage = Math.max(lastPercentage, Math.min(100, percentage));
    reportProgress?.({ percentage: lastPercentage, message });
  };
}

/**
 * Estimate how many bytes still need downloading before the tool can run with
 * the given options. Returns 0 when everything is already cached locally.
 */
export async function estimateUncachedDownloadBytes(
  options: BackgroundRemoverOptions,
  origin: string
): Promise<number> {
  const base = publicPath(origin);
  const resources = await loadResourceMap(base);
  if (!resources) return 0;

  let bytes = 0;
  for (const key of resourceKeys(options)) {
    const entry = resources[key];
    if (!entry?.chunks?.length) continue;
    // Chunks of a resource are fetched together, so the first chunk is a
    // reliable proxy for whether the whole resource is cached.
    const firstChunkUrl = new URL(entry.chunks[0].name, base).toString();
    if (!(await isImglyAssetCached(firstChunkUrl))) {
      bytes += entry.size;
    }
  }
  return bytes;
}

export async function removeImageBackground(
  fileName: string,
  imageData: Blob | ArrayBuffer | Uint8Array,
  options: BackgroundRemoverOptions,
  origin: string,
  reportProgress?: (progress: ToolProgress) => void
): Promise<ToolResult> {
  // Back /imgly/ fetches with Cache Storage so the model persists between runs
  // regardless of server cache headers.
  installImglyCache();

  const base = publicPath(origin);
  const resources = await loadResourceMap(base);
  const expectedSizes = new Map(
    resourceKeys(options).flatMap((key) => {
      const size = resources?.[key]?.size;
      return size === undefined ? [] : [[key, size] as const];
    })
  );
  const progress = createBackgroundRemovalProgress(
    expectedSizes,
    reportProgress
  );

  const source =
    imageData instanceof Blob
      ? imageData
      : new Blob([
          (imageData instanceof Uint8Array
            ? imageData
            : new Uint8Array(imageData)) as unknown as BlobPart,
        ]);

  const resultBlob = await removeBackground(source, {
    publicPath: base,
    model: MODEL_BY_QUALITY[options.quality],
    device: options.useGpu ? "gpu" : "cpu",
    // We already run inside a Web Worker; avoid nested worker proxying.
    proxyToWorker: false,
    output: { format: "image/png" },
    progress,
  });

  const outName = fileName.replace(/\.[^./\\]+$/, "") + "_no-bg.png";

  return {
    type: "files",
    files: [
      {
        name: outName,
        mimeType: "image/png",
        blob: resultBlob,
        sizeAfter: resultBlob.size,
      },
    ],
  };
}
