/**
 * A TypeScript client for this site's generation API (the OpenAPI file is
 * at /openapi.json). One file, no dependencies: copy it into your project.
 *
 * On a server, pass an API key from the API keys page:
 *
 *   const client = createClient({ baseUrl: 'https://example.com', apiKey });
 *   const generation = await client.create(
 *     { model: 'model-id', input: { prompt: 'A lighthouse at dawn.' } },
 *     { idempotencyKey: 'order-1234' }
 *   );
 *   const done = await client.wait(generation.id);
 *
 * Keep the key on your server, never in a browser or an app: anyone who
 * reads it can spend your credits. In a browser on this site, leave out
 * `apiKey` and `baseUrl`: requests go to the same site with the visitor's
 * own session.
 *
 * Creating sends an `Idempotency-Key`: your own id for the request (best),
 * or a new random one. A broken connection, a 5xx, a 429 or a 408 is sent
 * again with the same key and the same body, so a lost answer never charges
 * twice; any other refusal is answered at once. When the retries run out,
 * the error carries the key: send the request again later with it.
 */

export type GenerationStatus =
  | 'queued'
  | 'processing'
  | 'succeeded'
  | 'partially_succeeded'
  | 'failed'
  | 'canceled';

export type GenerationOutput = {
  id: string;
  index: number;
  /** `expired`: past the site's retention period, the file is deleted. */
  status: 'pending' | 'delivered' | 'failed' | 'withheld' | 'expired';
  url: string | null;
  content_type: string | null;
  bytes: number | null;
  duration_seconds: number | null;
  width: number | null;
  height: number | null;
};

export type Generation = {
  id: string;
  object: 'generation';
  model: string;
  status: GenerationStatus;
  created_at: string;
  finished_at: string | null;
  outputs: GenerationOutput[];
  billing: {
    mode: 'full' | 'trial';
    reserved_credits: number;
    charged_credits: number | null;
    refunded_credits: number | null;
  };
  review: { pending: true; since: string } | null;
  held_credits: number;
  error: { code: string; message: string } | null;
};

export type GenerationList = {
  object: 'list';
  items: Generation[];
  next_cursor: string | null;
};

export type CreateGenerationRequest = {
  /** A model id from the OpenAPI file's `x-models`. */
  model: string;
  /** The model's own input fields. */
  input: Record<string, unknown>;
  /** The most this request may reserve; it only lowers the limits. */
  max_credits?: number;
};

/**
 * An error the API answered with, by its code (the OpenAPI file's
 * `x-error-codes`), or one of the client's own, with `status` 0:
 * `network_error`, `timeout` and `not_ready` (an output without a file).
 */
export class OpenApiError extends Error {
  readonly status: number;
  readonly code: string;
  /** Seconds the API asked to wait before trying again. */
  readonly retryAfterSeconds?: number;
  /** The key a create was sent with: send it again with this key. */
  readonly idempotencyKey?: string;
  /** Fields of the error besides its code and message. */
  readonly details?: Record<string, unknown>;

  constructor(init: {
    status: number;
    code: string;
    message: string;
    retryAfterSeconds?: number;
    idempotencyKey?: string;
    details?: Record<string, unknown>;
  }) {
    super(init.message);
    this.name = 'OpenApiError';
    this.status = init.status;
    this.code = init.code;
    this.retryAfterSeconds = init.retryAfterSeconds;
    this.idempotencyKey = init.idempotencyKey;
    this.details = init.details;
  }
}

/** Where a browser keeps a create's key until it is answered (localStorage). */
export type KeyValueStore = {
  getItem(key: string): string | null;
  setItem(key: string, value: string): void;
  removeItem(key: string): void;
};

export type ClientOptions = {
  /** The site's address, such as https://example.com; leave out in a browser on it. */
  baseUrl?: string;
  /** An API key, on a server only. */
  apiKey?: string;
  fetch?: (url: string, init?: RequestInit) => Promise<Response>;
  /** Time for each request before it counts as a broken connection. */
  timeoutMs?: number;
  /** Times a request is sent again. */
  maxRetries?: number;
  /** A longer `Retry-After` is not waited out: the error goes to the caller. */
  maxRetryAfterSeconds?: number;
  sleep?: (ms: number, signal?: AbortSignal) => Promise<void>;
  /** Keeps a create's key across a page reload, by its `requestId`. */
  store?: KeyValueStore;
};

export type CreateOptions = {
  /** Your own id for this request; a new random one by default. */
  idempotencyKey?: string;
  /** With `store`: the same request after a reload keeps its first key. */
  requestId?: string;
  signal?: AbortSignal;
  headers?: Record<string, string>;
};

export type WaitOptions = {
  signal?: AbortSignal;
  /** Every reading of the generation, the last one included. */
  onUpdate?: (generation: Generation) => void;
  /**
   * How long to go on waiting, readings included: no reading or pause starts
   * past it, and a reading gets only the time left, so waiting stops with a
   * `timeout` error after about this long. It limits the waiting and the
   * requests, never a result in hand: a reading that shows the generation
   * ended is returned, even one that took the last of the time.
   */
  timeoutMs?: number;
};

export type GenerationClient = {
  create(
    request: CreateGenerationRequest,
    options?: CreateOptions
  ): Promise<Generation>;
  get(id: string, options?: { signal?: AbortSignal }): Promise<Generation>;
  list(options?: {
    limit?: number;
    cursor?: string;
    module?: 'image' | 'video' | 'audio';
    signal?: AbortSignal;
  }): Promise<GenerationList>;
  cancel(id: string, options?: { signal?: AbortSignal }): Promise<Generation>;
  /** An output's file; signed links open without the key. */
  download(
    output: { url: string | null } | string,
    options?: { signal?: AbortSignal; range?: string }
  ): Promise<Response>;
  /**
   * Reads the generation until it ends: 5 seconds apart, then up to 10, or
   * as a 429's or 503's `Retry-After` says. With `timeoutMs` it stops
   * waiting at that time; a generation it read as ended is still returned.
   */
  wait(id: string, options?: WaitOptions): Promise<Generation>;
};

const ENDED = new Set<GenerationStatus>([
  'succeeded',
  'partially_succeeded',
  'failed',
  'canceled',
]);
const FIRST_WAIT_MS = 5_000;
const LONGEST_WAIT_MS = 10_000;

function defaultSleep(ms: number, signal?: AbortSignal) {
  return new Promise<void>((resolve, reject) => {
    if (signal?.aborted) return reject(signal.reason);
    const timer = setTimeout(() => {
      signal?.removeEventListener('abort', stop);
      resolve();
    }, ms);
    const stop = () => {
      clearTimeout(timer);
      reject(signal!.reason);
    };
    signal?.addEventListener('abort', stop, { once: true });
  });
}

/** Seconds from a `Retry-After` header: a number of seconds or a date. */
function retryAfter(response: Response): number | undefined {
  const value = response.headers.get('retry-after');
  if (value === null) return undefined;
  if (/^\d+$/.test(value.trim())) return Number(value);
  const at = Date.parse(value);
  return Number.isNaN(at)
    ? undefined
    : Math.max(0, Math.ceil((at - Date.now()) / 1000));
}

async function answeredError(response: Response, idempotencyKey?: string) {
  let body: unknown;
  try {
    body = await response.json();
  } catch {
    body = undefined;
  }
  const error = (body as { error?: Record<string, unknown> } | undefined)
    ?.error;
  const { code, message, ...details } = error ?? {};
  return new OpenApiError({
    status: response.status,
    code:
      typeof code === 'string'
        ? code
        : response.status >= 500
          ? 'temporarily_unavailable'
          : 'request_failed',
    message: typeof message === 'string' ? message : `HTTP ${response.status}`,
    retryAfterSeconds: retryAfter(response),
    idempotencyKey,
    details: Object.keys(details).length ? details : undefined,
  });
}

/** Statuses worth sending the same request again for. */
const retryable = (status: number) =>
  status === 0 || status === 408 || status === 429 || status >= 500;

/**
 * `work`'s outcome, or `signal`'s reason as soon as it aborts: a `fetch` of
 * the caller's own, or the body it answered with, that does not stop when
 * aborted cannot hold a request past its time.
 */
function untilAborted<T>(work: Promise<T>, signal: AbortSignal) {
  return new Promise<T>((resolve, reject) => {
    const stop = () => reject(signal.reason);
    signal.addEventListener('abort', stop, { once: true });
    if (signal.aborted) stop();
    work
      .then(resolve, reject)
      .finally(() => signal.removeEventListener('abort', stop));
  });
}

export function createClient(options: ClientOptions = {}): GenerationClient {
  const base = (options.baseUrl ?? '').replace(/\/+$/, '');
  const apiKey = options.apiKey;
  const send = options.fetch ?? ((url, init) => fetch(url, init));
  const timeoutMs = options.timeoutMs ?? 30_000;
  const maxRetries = options.maxRetries ?? 5;
  const maxRetryAfter = options.maxRetryAfterSeconds ?? 60;
  const sleep = options.sleep ?? defaultSleep;
  const store = options.store;
  const auth = (): Record<string, string> =>
    apiKey ? { Authorization: `Bearer ${apiKey}` } : {};

  /**
   * One request; a broken connection or a timeout is a `status` 0 error.
   * The answer is read whole under the same time limit and the caller's
   * signal, so a body that stalls or is cut off counts as a broken
   * connection too, and so does a success whose JSON does not parse: it was
   * never really received. `limitMs` shortens the time limit. A file
   * (`stream`) has no time limit and stays tied to the caller's signal while
   * its body is read.
   */
  async function request(
    url: string,
    init: RequestInit,
    signal: AbortSignal | undefined,
    {
      idempotencyKey,
      stream = false,
      limitMs,
    }: { idempotencyKey?: string; stream?: boolean; limitMs?: number }
  ): Promise<Response> {
    if (signal?.aborted) throw signal.reason;
    const controller = new AbortController();
    const stop = () => controller.abort(signal?.reason);
    let timedOut = false;
    let timer: ReturnType<typeof setTimeout> | undefined;
    if (!stream) {
      signal?.addEventListener('abort', stop, { once: true });
      timer = setTimeout(
        () => {
          timedOut = true;
          controller.abort();
        },
        Math.min(timeoutMs, limitMs ?? timeoutMs)
      );
    }
    try {
      if (stream) return await send(url, { ...init, signal });
      const response = await untilAborted(
        send(url, { ...init, signal: controller.signal }),
        controller.signal
      );
      const body = await untilAborted(response.text(), controller.signal);
      if (response.ok) JSON.parse(body);
      // An empty body stays none (a 304 may have no other).
      return new Response(body === '' ? null : body, {
        status: response.status,
        statusText: response.statusText,
        headers: response.headers,
      });
    } catch (error) {
      if (signal?.aborted) throw signal.reason;
      throw new OpenApiError({
        status: 0,
        code: timedOut ? 'timeout' : 'network_error',
        message: timedOut
          ? 'The request took too long.'
          : `The request or its answer did not get through: ${error instanceof Error ? error.message : String(error)}`,
        idempotencyKey,
      });
    } finally {
      clearTimeout(timer);
      signal?.removeEventListener('abort', stop);
    }
  }

  /**
   * Sends a request until it is answered for good: a broken connection, a
   * 5xx, a 429 or a 408 again after `Retry-After` or a growing, jittered
   * wait. Returns the answer, or throws the last error.
   */
  async function attempt(
    url: string,
    init: RequestInit,
    signal: AbortSignal | undefined,
    {
      retry = true,
      idempotencyKey,
      stream,
      limitMs,
    }: {
      retry?: boolean;
      idempotencyKey?: string;
      stream?: boolean;
      limitMs?: number;
    }
  ): Promise<Response> {
    for (let tries = 0; ; tries++) {
      let error: OpenApiError;
      try {
        const response = await request(url, init, signal, {
          idempotencyKey,
          stream,
          limitMs,
        });
        if (response.ok) return response;
        error = await answeredError(response, idempotencyKey);
      } catch (caught) {
        if (!(caught instanceof OpenApiError)) throw caught;
        error = caught;
      }
      const wait = error.retryAfterSeconds;
      if (
        !retry ||
        !retryable(error.status) ||
        tries >= maxRetries ||
        (wait !== undefined && wait > maxRetryAfter)
      )
        throw error;
      await sleep(
        wait !== undefined
          ? wait * 1000
          : Math.min(30_000, 1000 * 2 ** tries) * (0.5 + Math.random() / 2),
        signal
      );
    }
  }

  const json = async <T>(response: Response) => (await response.json()) as T;
  const path = (id: string) =>
    `${base}/api/v1/generations/${encodeURIComponent(id)}`;

  async function create(
    body: CreateGenerationRequest,
    { idempotencyKey, requestId, signal, headers }: CreateOptions = {}
  ) {
    const text = JSON.stringify(body);
    const slot =
      store && requestId !== undefined
        ? `open-api-request:${requestId}`
        : undefined;
    let key = idempotencyKey;
    if (!key && slot) {
      try {
        const saved = JSON.parse(store!.getItem(slot) ?? 'null');
        if (saved?.body === text && typeof saved.key === 'string')
          key = saved.key;
      } catch {
        // A damaged entry is a new request.
      }
    }
    key ??= crypto.randomUUID();
    if (slot) store!.setItem(slot, JSON.stringify({ key, body: text }));
    try {
      const response = await attempt(
        `${base}/api/v1/generations`,
        {
          method: 'POST',
          headers: {
            ...headers,
            ...auth(),
            'Content-Type': 'application/json',
            'Idempotency-Key': key,
          },
          body: text,
        },
        signal,
        { idempotencyKey: key }
      );
      // The key is forgotten only once the whole answer is read.
      const generation = await json<Generation>(response);
      if (slot) store!.removeItem(slot);
      return generation;
    } catch (error) {
      // A refusal is an answer; a request still unanswered keeps its key.
      if (slot && error instanceof OpenApiError && !retryable(error.status))
        store!.removeItem(slot);
      throw error;
    }
  }

  async function get(id: string, { signal }: { signal?: AbortSignal } = {}) {
    return json<Generation>(
      await attempt(path(id), { headers: auth() }, signal, {})
    );
  }

  async function list({
    limit,
    cursor,
    module,
    signal,
  }: Parameters<GenerationClient['list']>[0] = {}) {
    const query = new URLSearchParams();
    if (limit !== undefined) query.set('limit', String(limit));
    if (cursor !== undefined) query.set('cursor', cursor);
    if (module !== undefined) query.set('module', module);
    const search = query.toString();
    return json<GenerationList>(
      await attempt(
        `${base}/api/v1/generations${search ? `?${search}` : ''}`,
        { headers: auth() },
        signal,
        {}
      )
    );
  }

  async function cancel(id: string, { signal }: { signal?: AbortSignal } = {}) {
    return json<Generation>(
      await attempt(
        `${path(id)}/cancel`,
        { method: 'POST', headers: auth() },
        signal,
        { retry: false }
      )
    );
  }

  async function download(
    output: { url: string | null } | string,
    { signal, range }: { signal?: AbortSignal; range?: string } = {}
  ) {
    const url = typeof output === 'string' ? output : output.url;
    if (!url)
      throw new OpenApiError({
        status: 0,
        code: 'not_ready',
        message: 'This output has no file yet.',
      });
    // The key goes to this site's own unsigned links only.
    let ownLink = false;
    if (apiKey && base) {
      const target = new URL(url, base);
      ownLink =
        target.origin === new URL(base).origin &&
        !target.searchParams.has('sig');
    }
    return attempt(
      url,
      {
        headers: {
          ...(ownLink ? auth() : {}),
          ...(range ? { Range: range } : {}),
        },
      },
      signal,
      { retry: false, stream: true }
    );
  }

  async function wait(
    id: string,
    { signal, onUpdate, timeoutMs: limit }: WaitOptions = {}
  ) {
    let interval = FIRST_WAIT_MS;
    // The time used so far, from the start: everything as it really took
    // (readings, the callback, pauses that ran late), and each pause at least
    // as long as it was asked (a `sleep` of the caller's own may return
    // sooner).
    const start = Date.now();
    let unslept = 0;
    const spent = () => Date.now() - start + unslept;
    const timedOut = () =>
      new OpenApiError({
        status: 0,
        code: 'timeout',
        message: 'The generation did not end in time.',
      });
    // A pause that would end past the time is not taken: checked after the
    // callback, whose own time counts.
    const pause = async (ms: number) => {
      if (limit !== undefined && spent() + ms > limit) throw timedOut();
      const before = Date.now();
      await sleep(ms, signal);
      unslept += Math.max(0, ms - (Date.now() - before));
    };
    const longer = () => {
      const now = interval;
      interval = Math.min(interval + 1000, LONGEST_WAIT_MS);
      return now;
    };
    for (;;) {
      // Nothing is sent once the time is up (a pause ran late).
      if (limit !== undefined && spent() >= limit) throw timedOut();
      let generation: Generation;
      try {
        generation = await json<Generation>(
          await attempt(path(id), { headers: auth() }, signal, {
            retry: false,
            // A reading never runs past the time left.
            limitMs:
              limit === undefined ? undefined : Math.max(1, limit - spent()),
          })
        );
      } catch (error) {
        if (!(error instanceof OpenApiError)) throw error;
        // A 429 or a 503 says how long to wait; past a minute, it is the
        // caller's to decide, as for any request.
        const asked = error.retryAfterSeconds;
        if (
          !retryable(error.status) ||
          (asked !== undefined && asked > maxRetryAfter)
        )
          throw error;
        await pause(asked !== undefined ? asked * 1000 : longer());
        continue;
      }
      onUpdate?.(generation);
      // An ended generation is the answer, whatever its callback took.
      if (ENDED.has(generation.status)) return generation;
      await pause(longer());
    }
  }

  return { create, get, list, cancel, download, wait };
}