// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

import { authFetch } from "@/features/auth";
import { formatFastApiDetail } from "@/lib/format-fastapi-error";
import { hubTokenHeader } from "../lib/hub-token-header";
import { abortError, withAbort } from "../lib/abort-signals";
import type { TransportMode } from "./constants";

function parseErrorText(status: number, body: unknown): string {
  if (body && typeof body === "object") {
    const detail = (body as { detail?: unknown }).detail;
    const formatted = formatFastApiDetail(detail);
    if (status === 405) {
      return `${formatted || "Method Not Allowed"} - the Unsloth backend did not accept this API method. Restart Unsloth so the frontend and backend are on the same build.`;
    }
    if (formatted) return formatted;
    const message = (body as { message?: unknown }).message;
    if (typeof message === "string" && message) return message;
  }
  if (status === 405) {
    return "Method Not Allowed - the Unsloth backend did not accept this API method. Restart Unsloth so the frontend and backend are on the same build.";
  }
  return `Request failed (${status})`;
}

async function parseJsonOrThrow<T>(response: Response): Promise<T> {
  const body = await response.json().catch(() => null);
  if (!response.ok) {
    throw new Error(parseErrorText(response.status, body));
  }
  return body as T;
}

export type DownloadJobState =
  | "idle"
  | "running"
  | "complete"
  | "error"
  | "cancelled"
  | "cancelling";

export interface DownloadJobStatus {
  state: DownloadJobState;
  error?: string | null;
  generation?: number;
}

// "repository_owned": a dictation model download holds this repository's cache.
export type DownloadStartState =
  | DownloadJobState
  | "deleting"
  | "repository_owned";

export interface DownloadStartResult {
  state: DownloadStartState;
  accepted: boolean;
  generation?: number;
  // Present only when the start adopted a job another client had already
  // begun: the transport it is really running on.
  transport?: TransportMode | null;
  // And its cancel marker, when that job had fallen back from Xet to HTTP.
  cancel_transport?: TransportMode | null;
}

export interface ActiveModelDownload {
  repo_id?: string;
  variant: string | null;
  transport?: TransportMode | null;
  // Set only on a Xet run that fell back to HTTP: stopping it still leaves a
  // restart-only partial, so this and not `transport` decides the stop control.
  cancel_transport?: TransportMode | null;
  state: DownloadJobState;
  generation?: number;
  // Scoped jobs only: the exact files this job is fetching. Every file set of one repo rides the same "@scope" slot, so an
  // adopting client needs it to tell its own transfer from a sibling's. Absent from an older backend means unprovable.
  files?: string[] | null;
}

export interface ActiveDatasetDownload {
  repo_id: string;
  variant: null;
  transport?: TransportMode | null;
  cancel_transport?: TransportMode | null;
  state: DownloadJobState;
  generation?: number;
}

type ActiveModelDownloadsOptions = {
  fresh?: boolean;
};

const ACTIVE_MODEL_DOWNLOADS_DEDUPE_MS = 750;
const ACTIVE_MODEL_DOWNLOADS_CACHE_MAX = 200;
const activeModelDownloadsInFlight = new Map<
  string,
  Promise<readonly ActiveModelDownload[]>
>();
const activeModelDownloadsCache = new Map<
  string,
  { expiresAt: number; downloads: readonly ActiveModelDownload[] }
>();

function pruneActiveModelDownloadsCache(now = Date.now()): void {
  for (const [key, entry] of activeModelDownloadsCache) {
    if (entry.expiresAt <= now) activeModelDownloadsCache.delete(key);
  }
  while (activeModelDownloadsCache.size > ACTIVE_MODEL_DOWNLOADS_CACHE_MAX) {
    const oldest = activeModelDownloadsCache.keys().next().value;
    if (oldest === undefined) break;
    activeModelDownloadsCache.delete(oldest);
  }
}

export type TransportMarker = TransportMode | null;

export interface TransportStatus {
  has_partial: boolean;
  last_transport: TransportMarker;
  resumable: boolean;
}

export interface DownloadProgressResponse {
  downloaded_bytes: number;
  completed_bytes: number;
  complete_on_disk: boolean;
  expected_bytes: number;
  progress: number;
  cache_path: string | null;
  /**
   * Whether the backend found anything for THIS target (variant), as opposed to the shared
   * repo cache directory existing. Null where it cannot say, and absent from an older
   * backend -- both of which leave the repo-level rule in charge.
   */
  target_present?: boolean | null;
  /** False when the cache could not be scanned at all: unknown, not empty. Absent from an
   *  older backend, which is also unknown. */
  cache_measured?: boolean;
}

export type DownloadProgressOptions = {
  signal?: AbortSignal;
  hfToken?: string | null;
  expectedBytes?: number;
};

export type GgufDownloadProgressOptions = DownloadProgressOptions & {
  variant: string;
};

export type {
  DownloadTransportCapability,
  DownloadTransportCapabilities,
} from "./transport-capabilities";
import {
  DOWNLOAD_TRANSPORT_CAPABILITIES_FALLBACK,
  normalizeDownloadTransportCapabilities,
} from "./transport-capabilities";
import type { DownloadTransportCapabilities } from "./transport-capabilities";

// The Auto verdict is dynamic: a stalled transfer demotes this machine mid-session, so an
// indefinite cache would keep serving the pre-stall answer until the page reloads. Refreshing is
// cheap because the endpoint answers from cached health without a network probe.
const DOWNLOAD_TRANSPORT_CAPABILITIES_TTL_MS = 30_000;
let downloadTransportCapabilitiesCache: {
  expiresAt: number;
  probed: boolean;
  capabilities: DownloadTransportCapabilities;
} | null = null;
let downloadTransportCapabilitiesInFlight: Promise<DownloadTransportCapabilities> | null =
  null;

export async function getDownloadTransportCapabilities(options: {
  force?: boolean;
  // Reach for the Xet endpoint instead of answering from the cached verdict. Only the
  // download-start path sets this; the UI polls on render and must not connect per poll.
  probe?: boolean;
} = {}): Promise<DownloadTransportCapabilities> {
  const cached = downloadTransportCapabilitiesCache;
  // A cheap cached answer cannot satisfy a probe request, but a probed one satisfies everybody.
  const cacheUsable =
    cached !== null &&
    cached.expiresAt > Date.now() &&
    (!options.probe || cached.probed);
  if (!options.force && cacheUsable) {
    return cached.capabilities;
  }
  if (!options.force && !options.probe && downloadTransportCapabilitiesInFlight) {
    return downloadTransportCapabilitiesInFlight;
  }
  const request = authFetch(
    `/api/studio/download-transport-capabilities${options.probe ? "?probe=1" : ""}`,
  )
    .then(parseJsonOrThrow<unknown>)
    .then(normalizeDownloadTransportCapabilities)
    .then((capabilities) => {
      downloadTransportCapabilitiesCache = {
        expiresAt: Date.now() + DOWNLOAD_TRANSPORT_CAPABILITIES_TTL_MS,
        probed: options.probe === true,
        capabilities,
      };
      return capabilities;
    })
    .catch(() => DOWNLOAD_TRANSPORT_CAPABILITIES_FALLBACK)
    .finally(() => {
      if (downloadTransportCapabilitiesInFlight === request) {
        downloadTransportCapabilitiesInFlight = null;
      }
    });
  downloadTransportCapabilitiesInFlight = request;
  return request;
}

export async function startModelDownload(payload: {
  repo_id: string;
  gguf_variant?: string | null;
  hf_token?: string | null;
  use_xet?: boolean;
  // A partial-by-design download of `files` only (see DownloadRequest.scopeId).
  scope_id?: string | null;
  files?: string[];
  transport_mode?: "auto" | "xet" | "http";
}): Promise<DownloadStartResult & { job_key: string }> {
  const { hf_token, ...body } = payload;
  const headers = {
    "Content-Type": "application/json",
    ...hubTokenHeader(hf_token),
  };
  const response = await authFetch("/api/hub/download", {
    method: "POST",
    headers,
    body: JSON.stringify(body),
  });
  return parseJsonOrThrow(response);
}

export async function cancelModelDownload(payload: {
  repo_id: string;
  gguf_variant?: string | null;
  generation?: number | null;
  signal?: AbortSignal;
}): Promise<{ job_key: string; state: DownloadJobState }> {
  const { signal, ...body } = payload;
  const response = await authFetch("/api/hub/download/cancel", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(body),
    signal,
  });
  return parseJsonOrThrow(response);
}

export async function getModelDownloadStatus(
  repoId: string,
  ggufVariant?: string | null,
  signal?: AbortSignal,
): Promise<DownloadJobStatus> {
  const params = new URLSearchParams({ repo_id: repoId });
  if (ggufVariant) params.set("gguf_variant", ggufVariant);
  const response = await authFetch(`/api/hub/download-status?${params}`, {
    signal,
  });
  return parseJsonOrThrow<DownloadJobStatus>(response);
}

export async function getActiveModelDownloads(
  repoId: string,
  signal?: AbortSignal,
  options: ActiveModelDownloadsOptions = {},
): Promise<readonly ActiveModelDownload[]> {
  return getActiveModelDownloadsForKey(repoId, signal, options);
}

export async function getAllActiveModelDownloads(
  signal?: AbortSignal,
  options: ActiveModelDownloadsOptions = {},
): Promise<readonly ActiveModelDownload[]> {
  return getActiveModelDownloadsForKey(null, signal, options);
}

async function getActiveModelDownloadsForKey(
  repoId: string | null,
  signal?: AbortSignal,
  options: ActiveModelDownloadsOptions = {},
): Promise<readonly ActiveModelDownload[]> {
  const trimmedRepoId = repoId?.trim() ?? "";
  const key = trimmedRepoId ? trimmedRepoId.toLowerCase() : "*";
  if (signal?.aborted) return Promise.reject(abortError(signal));
  const useCachedResult = !options.fresh;
  const cached = useCachedResult ? activeModelDownloadsCache.get(key) : undefined;
  const now = Date.now();
  if (cached && cached.expiresAt > now) {
    return withAbort(Promise.resolve(cached.downloads), signal);
  }

  let request = useCachedResult ? activeModelDownloadsInFlight.get(key) : undefined;
  if (!request) {
    const fetchRequest = (async () => {
      const params = new URLSearchParams();
      if (trimmedRepoId) params.set("repo_id", trimmedRepoId);
      const query = params.toString();
      const response = await authFetch(
        `/api/hub/active-downloads${query ? `?${query}` : ""}`,
        options.fresh && signal ? { signal } : undefined,
      );
      const data = await parseJsonOrThrow<{
        downloads: ActiveModelDownload[];
      }>(response);
      activeModelDownloadsCache.set(key, {
        expiresAt: Date.now() + ACTIVE_MODEL_DOWNLOADS_DEDUPE_MS,
        downloads: data.downloads,
      });
      pruneActiveModelDownloadsCache();
      return data.downloads;
    })();
    request = useCachedResult
      ? fetchRequest.finally(() => {
          activeModelDownloadsInFlight.delete(key);
        })
      : fetchRequest;
    if (useCachedResult) {
      activeModelDownloadsInFlight.set(key, request);
    }
  }
  return withAbort(request, signal);
}

export async function getActiveDatasetDownloads(
  signal?: AbortSignal,
  repoId?: string,
): Promise<ActiveDatasetDownload[]> {
  // No repo id lists every active dataset download, which is what hydration
  // wants; one narrows it, which is what a single repo view wants.
  const params = repoId ? `?${new URLSearchParams({ repo_id: repoId })}` : "";
  const response = await authFetch(`/api/hub/datasets/active-downloads${params}`, {
    signal,
  });
  const data = await parseJsonOrThrow<{
    downloads: ActiveDatasetDownload[];
  }>(response);
  return data.downloads;
}

export async function startDatasetDownload(payload: {
  repo_id: string;
  hf_token?: string | null;
  use_xet?: boolean;
  transport_mode?: "auto" | "xet" | "http";
}): Promise<DownloadStartResult & { repo_id: string }> {
  const { hf_token, ...body } = payload;
  const headers = {
    "Content-Type": "application/json",
    ...hubTokenHeader(hf_token),
  };
  const response = await authFetch("/api/hub/datasets/download", {
    method: "POST",
    headers,
    body: JSON.stringify(body),
  });
  return parseJsonOrThrow(response);
}

export async function cancelDatasetDownload(payload: {
  repo_id: string;
  generation?: number | null;
  signal?: AbortSignal;
}): Promise<{ repo_id: string; state: DownloadJobState }> {
  const { signal, ...body } = payload;
  const response = await authFetch("/api/hub/datasets/download/cancel", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(body),
    signal,
  });
  return parseJsonOrThrow(response);
}

export async function getDatasetDownloadStatus(
  repoId: string,
  signal?: AbortSignal,
): Promise<DownloadJobStatus> {
  const params = new URLSearchParams({ repo_id: repoId });
  const response = await authFetch(`/api/hub/datasets/download-status?${params}`, {
    signal,
  });
  return parseJsonOrThrow<DownloadJobStatus>(response);
}

export async function getModelTransportStatus(
  repoId: string,
  ggufVariant?: string | null,
  hfToken?: string | null,
  signal?: AbortSignal,
): Promise<TransportStatus> {
  const params = new URLSearchParams({ repo_id: repoId });
  if (ggufVariant) params.set("gguf_variant", ggufVariant);
  const response = await authFetch(`/api/hub/transport-status?${params}`, {
    headers: hubTokenHeader(hfToken),
    signal,
  });
  return parseJsonOrThrow<TransportStatus>(response);
}

export async function getDatasetTransportStatus(
  repoId: string,
  signal?: AbortSignal,
): Promise<TransportStatus> {
  const params = new URLSearchParams({ repo_id: repoId });
  const response = await authFetch(`/api/hub/datasets/transport-status?${params}`, {
    signal,
  });
  return parseJsonOrThrow<TransportStatus>(response);
}

export async function getGgufDownloadProgress(
  repoId: string,
  options: GgufDownloadProgressOptions,
): Promise<DownloadProgressResponse> {
  const { expectedBytes = 0, hfToken, signal, variant } = options;
  const params = new URLSearchParams({
    repo_id: repoId,
    variant,
    expected_bytes: String(expectedBytes),
  });
  const response = await authFetch(`/api/hub/gguf-download-progress?${params}`, {
    headers: hubTokenHeader(hfToken),
    signal,
  });
  return parseJsonOrThrow(response);
}

export async function getDownloadProgress(
  repoId: string,
  options: DownloadProgressOptions,
): Promise<DownloadProgressResponse> {
  const { expectedBytes = 0, hfToken, signal } = options;
  const params = new URLSearchParams({ repo_id: repoId });
  if (expectedBytes > 0) params.set("expected_bytes", String(expectedBytes));
  const response = await authFetch(`/api/hub/download-progress?${params}`, {
    headers: hubTokenHeader(hfToken),
    signal,
  });
  return parseJsonOrThrow(response);
}

export async function getDatasetDownloadProgress(
  repoId: string,
  options: DownloadProgressOptions,
): Promise<DownloadProgressResponse> {
  const { expectedBytes = 0, hfToken, signal } = options;
  const params = new URLSearchParams({ repo_id: repoId });
  if (expectedBytes > 0) params.set("expected_bytes", String(expectedBytes));
  const response = await authFetch(`/api/hub/datasets/download-progress?${params}`, {
    headers: hubTokenHeader(hfToken),
    signal,
  });
  return parseJsonOrThrow(response);
}
