/**
 * A typed client for the Prime Intellect compute API.
 *
 * Written by hand because no first-party TypeScript SDK exists — the official
 * SDK and CLI are Python. The surface here is deliberately narrow: PIG reads
 * GPU availability to populate its inventory and does nothing else. It cannot
 * provision, terminate, or spend money, and the API key it holds should be
 * scoped so that it could not even if the code tried.
 *
 * Two operational realities shape this file.
 *
 * **Rate limits are undocumented.** There is no published quota, so the client
 * cannot pace itself against a known budget. It instead backs off empirically:
 * exponential with jitter on 429 and 5xx, honouring `Retry-After` when the
 * server sends one, and giving up rather than hammering.
 *
 * **The upstream shape may drift.** Responses are parsed defensively; unknown
 * fields are preserved verbatim in `raw` rather than dropped, so a field that
 * appears upstream tomorrow is already captured today.
 */

export interface PrimeClientOptions {
  apiKey: string;
  baseUrl?: string;
  /** Total attempts per request, including the first. */
  maxAttempts?: number;
  /** Ceiling on backoff between attempts. */
  maxBackoffMs?: number;
  /** Per-request timeout. */
  timeoutMs?: number;
  fetchImpl?: typeof fetch;
  onRetry?: (info: { attempt: number; delayMs: number; reason: string }) => void;
}

export class PrimeApiError extends Error {
  constructor(
    message: string,
    readonly status: number,
    readonly body?: string,
  ) {
    super(message);
    this.name = 'PrimeApiError';
  }

  /** Retrying a 4xx that is not 429 will fail identically every time. */
  get isRetryable(): boolean {
    return this.status === 429 || this.status >= 500;
  }
}

/** A GPU availability listing, as returned by the availability endpoints. */
export interface PrimeGpuListing {
  cloudId?: string;
  gpuType?: string;
  socket?: string;
  provider?: string;
  region?: string;
  dataCenter?: string;
  country?: string;
  gpuCount?: number;
  gpuMemory?: number;
  vcpu?: { defaultCount?: number } | number;
  memory?: { defaultCount?: number } | number;
  disk?: {
    minCount?: number;
    defaultCount?: number;
    maxCount?: number;
    pricePerUnit?: number;
  };
  internetSpeed?: number;
  interconnect?: number;
  interconnectType?: string;
  provisioningTime?: number;
  stockStatus?: string;
  security?: string;
  prices?: {
    onDemand?: number | null;
    communityPrice?: number | null;
    isVariable?: boolean;
    currency?: string;
  };
  images?: string[];
  isSpot?: boolean;
  prepaidTime?: number;
  /** Everything the upstream sent, including fields not modelled above. */
  raw: Record<string, unknown>;
}

export interface PrimeAvailabilityQuery {
  regions?: string[];
  gpuCount?: number;
  gpuType?: string;
  socket?: string;
  security?: 'secure_cloud' | 'community_cloud';
  dataCenterId?: string;
  cloudId?: string;
  page?: number;
  /** Upstream caps this at 100. Values above are clamped rather than rejected. */
  pageSize?: number;
}

const DEFAULT_BASE_URL = 'https://api.primeintellect.ai';

export class PrimeClient {
  private readonly apiKey: string;
  private readonly baseUrl: string;
  private readonly maxAttempts: number;
  private readonly maxBackoffMs: number;
  private readonly timeoutMs: number;
  private readonly fetchImpl: typeof fetch;
  private readonly onRetry: PrimeClientOptions['onRetry'];

  constructor(options: PrimeClientOptions) {
    if (!options.apiKey) throw new Error('PrimeClient requires an apiKey.');
    this.apiKey = options.apiKey;
    this.baseUrl = (options.baseUrl ?? DEFAULT_BASE_URL).replace(/\/+$/, '');
    this.maxAttempts = options.maxAttempts ?? 5;
    this.maxBackoffMs = options.maxBackoffMs ?? 30_000;
    this.timeoutMs = options.timeoutMs ?? 30_000;
    this.fetchImpl = options.fetchImpl ?? fetch;
    this.onRetry = options.onRetry;
  }

  /**
   * One page of GPU availability.
   *
   * Returns the listings plus `totalCount` so a caller can page without
   * guessing. Note that upstream paginates from page 1, not 0.
   */
  async listGpuAvailability(
    query: PrimeAvailabilityQuery = {},
  ): Promise<{ items: PrimeGpuListing[]; totalCount: number }> {
    const params = new URLSearchParams();
    for (const region of query.regions ?? []) params.append('regions', region);
    if (query.gpuCount != null) params.set('gpu_count', String(query.gpuCount));
    if (query.gpuType) params.set('gpu_type', query.gpuType);
    if (query.socket) params.set('socket', query.socket);
    if (query.security) params.set('security', query.security);
    if (query.dataCenterId) params.set('data_center_id', query.dataCenterId);
    if (query.cloudId) params.set('cloud_id', query.cloudId);
    params.set('page', String(query.page ?? 1));
    // Upstream rejects page_size above 100; clamp rather than surface a 400.
    params.set('page_size', String(Math.min(query.pageSize ?? 100, 100)));

    const body = await this.request<unknown>(`/api/v1/availability/gpus?${params}`);
    return normaliseListingPage(body);
  }

  /**
   * Every page of GPU availability, as an async iterable.
   *
   * Yields page by page rather than accumulating, so a large inventory does not
   * have to fit in memory at once and the caller can begin upserting
   * immediately. Stops when a page comes back empty, which also protects
   * against a `totalCount` that disagrees with reality.
   */
  async *iterateGpuAvailability(
    query: Omit<PrimeAvailabilityQuery, 'page'> = {},
  ): AsyncGenerator<PrimeGpuListing[], void, undefined> {
    let page = 1;
    let seen = 0;
    for (;;) {
      const { items, totalCount } = await this.listGpuAvailability({ ...query, page });
      if (items.length === 0) return;
      yield items;
      seen += items.length;
      if (totalCount > 0 && seen >= totalCount) return;
      page += 1;
      // A defensive ceiling. Without it a misbehaving upstream that always
      // returns a full page would loop until the process is killed.
      if (page > 1000) return;
    }
  }

  /** Multi-node cluster availability — the shape that can actually train. */
  async listMultiNodeAvailability(
    query: PrimeAvailabilityQuery = {},
  ): Promise<{ items: PrimeGpuListing[]; totalCount: number }> {
    const params = new URLSearchParams();
    if (query.gpuType) params.set('gpu_type', query.gpuType);
    if (query.gpuCount != null) params.set('gpu_count', String(query.gpuCount));
    params.set('page', String(query.page ?? 1));
    params.set('page_size', String(Math.min(query.pageSize ?? 100, 100)));

    const body = await this.request<unknown>(`/api/v1/availability/multi-node?${params}`);
    return normaliseListingPage(body);
  }

  /**
   * Cheap credential check. Used by the settings screen so an operator finds
   * out their key is wrong at configuration time rather than at 3am when the
   * sync silently stops.
   */
  async verifyCredentials(): Promise<{ ok: boolean; detail?: string }> {
    try {
      await this.listGpuAvailability({ pageSize: 1 });
      return { ok: true };
    } catch (error) {
      if (error instanceof PrimeApiError) {
        return {
          ok: false,
          detail:
            error.status === 401 || error.status === 403
              ? 'Key rejected. Check it has the Availability → Read scope and has not expired.'
              : `Upstream returned ${error.status}.`,
        };
      }
      return { ok: false, detail: error instanceof Error ? error.message : 'Unknown error' };
    }
  }

  private async request<T>(path: string): Promise<T> {
    let lastError: unknown;

    for (let attempt = 1; attempt <= this.maxAttempts; attempt++) {
      const controller = new AbortController();
      const timer = setTimeout(() => controller.abort(), this.timeoutMs);

      try {
        const response = await this.fetchImpl(`${this.baseUrl}${path}`, {
          headers: {
            authorization: `Bearer ${this.apiKey}`,
            accept: 'application/json',
            'user-agent': 'pig-crm/0.1 (+https://primeintellectgrowth.com)',
          },
          signal: controller.signal,
        });

        if (response.ok) return (await response.json()) as T;

        const text = await response.text().catch(() => '');
        const error = new PrimeApiError(
          `Prime Intellect API returned ${response.status}`,
          response.status,
          text.slice(0, 500),
        );
        if (!error.isRetryable || attempt === this.maxAttempts) throw error;

        // Honour Retry-After when offered; it is better information than any
        // backoff curve we could invent.
        const retryAfter = parseRetryAfter(response.headers.get('retry-after'));
        const delay = retryAfter ?? this.backoffMs(attempt);
        this.onRetry?.({ attempt, delayMs: delay, reason: `HTTP ${response.status}` });
        await sleep(delay);
        lastError = error;
      } catch (error) {
        if (error instanceof PrimeApiError) {
          if (!error.isRetryable || attempt === this.maxAttempts) throw error;
          lastError = error;
        } else {
          // Network failure or timeout. Both are worth retrying.
          if (attempt === this.maxAttempts) throw error;
          const delay = this.backoffMs(attempt);
          this.onRetry?.({
            attempt,
            delayMs: delay,
            reason: error instanceof Error ? error.message : 'network error',
          });
          await sleep(delay);
          lastError = error;
        }
      } finally {
        clearTimeout(timer);
      }
    }

    throw lastError ?? new Error('Prime Intellect request failed.');
  }

  /**
   * Exponential backoff with full jitter.
   *
   * Jitter matters more than the curve: without it, several workers that hit a
   * limit together retry together, and the thundering herd reproduces the
   * problem that caused the limit.
   */
  private backoffMs(attempt: number): number {
    const ceiling = Math.min(this.maxBackoffMs, 1000 * 2 ** (attempt - 1));
    return Math.round(Math.random() * ceiling);
  }
}

function parseRetryAfter(header: string | null): number | null {
  if (!header) return null;
  const seconds = Number(header);
  if (Number.isFinite(seconds)) return Math.max(0, seconds * 1000);
  const date = Date.parse(header);
  if (Number.isFinite(date)) return Math.max(0, date - Date.now());
  return null;
}

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

/**
 * Normalise a page of listings.
 *
 * The upstream has returned both `{items, totalCount}` and a bare array across
 * its endpoint generations, so both are accepted rather than assuming one and
 * breaking on the other.
 */
function normaliseListingPage(body: unknown): {
  items: PrimeGpuListing[];
  totalCount: number;
} {
  const rawItems: unknown[] = Array.isArray(body)
    ? body
    : Array.isArray((body as { items?: unknown[] })?.items)
      ? ((body as { items: unknown[] }).items ?? [])
      : [];

  const totalCount =
    typeof (body as { totalCount?: number })?.totalCount === 'number'
      ? (body as { totalCount: number }).totalCount
      : rawItems.length;

  return { items: rawItems.map(toListing), totalCount };
}

function toListing(raw: unknown): PrimeGpuListing {
  const r = (raw ?? {}) as Record<string, unknown>;
  return {
    cloudId: str(r.cloudId),
    gpuType: str(r.gpuType),
    socket: str(r.socket),
    provider: str(r.provider),
    region: str(r.region),
    dataCenter: str(r.dataCenter),
    country: str(r.country),
    gpuCount: num(r.gpuCount),
    gpuMemory: num(r.gpuMemory),
    vcpu: r.vcpu as PrimeGpuListing['vcpu'],
    memory: r.memory as PrimeGpuListing['memory'],
    disk: r.disk as PrimeGpuListing['disk'],
    internetSpeed: num(r.internetSpeed),
    interconnect: num(r.interconnect),
    interconnectType: str(r.interconnectType),
    provisioningTime: num(r.provisioningTime),
    stockStatus: str(r.stockStatus),
    security: str(r.security),
    prices: r.prices as PrimeGpuListing['prices'],
    images: Array.isArray(r.images) ? (r.images as string[]) : undefined,
    isSpot: typeof r.isSpot === 'boolean' ? r.isSpot : undefined,
    prepaidTime: num(r.prepaidTime),
    raw: r,
  };
}

const str = (v: unknown): string | undefined => (typeof v === 'string' ? v : undefined);
const num = (v: unknown): number | undefined => (typeof v === 'number' ? v : undefined);
