Jugar en el navegador

// Messages exchanged with `src/workers/decoder.worker.ts`. Both sides import this module, so a
// change to a payload breaks the build on the side that did not follow.
//
// The flow is always the same: `init` (download + session) answers with `progress`* and then
// `ready` or `error`; `logits` answers with `logits` or `error`; `dispose` answers with
// `disposed`. Every request carries an `id` the answer echoes, so the client can correlate
// replies even though the worker already serialises the work.

/** ONNX Runtime Web version; `scripts/copy-assets.mjs` copies that release into `public/ort/`. */
export const ORT_VERSION = '1.30.0';

/** Where the self-hosted ORT runtime (`.wasm` + loader) is served from. */
export const ORT_BASE = `/ort/${ORT_VERSION}/`;

/**
 * Cache API bucket for downloaded models. The suffix is bumped whenever what is stored (or how
 * it is keyed) changes, so an old browser cache is never read with new expectations.
 */
export const MODEL_CACHE = 'rukh-models-v1';

export type Backend = 'webgpu' | 'wasm';

export interface InitRequest {
  type: 'init';
  id: number;
  /** Registry stage id, only used for diagnostics and the cache log. */
  stage: string;
  /** Absolute or same-origin URL of the `.onnx` file. */
  url: string;
  /** Expected size in bytes; used for the progress bar when there is no `content-length`. */
  sizeBytes: number;
  /** Context window the stage was trained with (`DecoderConfig.block`), from the registry. */
  block: number;
  /** Vocabulary size the caller's tokenizer has; the model's output width must match it. */
  vocab: number;
}

export interface LogitsRequest {
  type: 'logits';
  id: number;
  /** Token ids of the prompt, already cropped to the context by the caller. */
  ids: number[];
}

export interface DisposeRequest {
  type: 'dispose';
  id: number;
}

export type WorkerRequest = InitRequest | LogitsRequest | DisposeRequest;

export interface ProgressMessage {
  type: 'progress';
  id: number;
  loaded: number;
  total: number;
}

export interface ReadyMessage {
  type: 'ready';
  id: number;
  backend: Backend;
  /** Why WebGPU was not used, when `backend` is `wasm` and a reason is known. */
  fallbackReason?: string;
  loadMs: number;
  /** The contract the session was checked against; `buildPrompt` crops to `block`. */
  block: number;
  vocab: number;
}

export interface LogitsMessage {
  type: 'logits';
  id: number;
  data: Float32Array;
  inferMs: number;
}

export interface DisposedMessage {
  type: 'disposed';
  id: number;
}

export interface ErrorMessage {
  type: 'error';
  id: number;
  message: string;
}

export type WorkerResponse =
  ProgressMessage | ReadyMessage | LogitsMessage | DisposedMessage | ErrorMessage;

const RESPONSE_TYPES = ['progress', 'ready', 'logits', 'disposed', 'error'] as const;

/**
 * Narrows whatever arrived through `postMessage` to a known response, or `null`. The worker is
 * ours, but the message channel is not typed at runtime and a stray message must not crash the
 * page.
 */
export function asResponse(data: unknown): WorkerResponse | null {
  if (typeof data !== 'object' || data === null) return null;
  const message = data as { type?: unknown; id?: unknown };
  if (typeof message.id !== 'number') return null;
  const type = message.type;
  if (typeof type !== 'string') return null;
  if (!(RESPONSE_TYPES as readonly string[]).includes(type)) return null;
  return data as WorkerResponse;
}

/** Percentage (0-100) of a progress message; `total` of 0 reports 0 rather than `NaN`. */
export function progressPercent(message: Pick<ProgressMessage, 'loaded' | 'total'>): number {
  if (message.total <= 0) return 0;
  return Math.min(100, Math.round((message.loaded / message.total) * 100));
}

/** Bytes as MB with one decimal, the unit the consent step and the progress bar speak. */
export function megabytes(bytes: number): number {
  return Math.round((bytes / 1_000_000) * 10) / 10;
}

// What the demo demands of a model file before it is allowed to play, and the checks that prove
// it. A downloaded `.onnx` is just bytes: nothing stops a stale cache entry, a half-renamed Hub
// file or a model trained with another vocabulary from loading cleanly and then answering with a
// logits vector that means something else. Every id the sampler reads out of that vector would be
// a different move, so the game would look plausible and be nonsense.
//
// Two things have to hold and both are checked:
//
//   * the **width** of the output equals the tokenizer's vocabulary (2030 UCI tokens). ORT Web
//     does not expose `metadata_props`, so the `rukh_vocab_size` the exporter writes into the file
//     is unreachable from the browser; what is reachable is the output shape ORT reports
//     (`outputMetadata`) and, failing that, the length of the first real answer;
//   * the **context** (`block`) the prompt is cropped to. Nothing in the session says it, so it
//     comes from the registry entry and travels back to the main thread in `ready`, which is what
//     keeps `buildPrompt` from hardcoding 200.
//
// The messages are in Spanish on purpose: they end up in the model panel, in front of a player.

/** The contract a loaded session is playing under. */
export interface ModelContract {
  /** Context window in tokens (`DecoderConfig.block`); the prompt is cropped to it. */
  block: number;
  /** Width of the logits vector, which must equal the tokenizer's vocabulary size. */
  vocab: number;
}

/** The part of ORT's `InferenceSession.ValueMetadata` this check needs. */
export interface OutputMetadata {
  name: string;
  isTensor?: boolean;
  /** Dimensions; a string entry is a symbolic axis (`batch`, `sequence`). */
  shape?: readonly (number | string)[];
}

/** The part of `InferenceSession` this check needs, so the tests can hand it a plain object. */
export interface SessionLike {
  readonly outputNames: readonly string[];
  readonly outputMetadata?: readonly OutputMetadata[];
}

/** Smallest context that can still hold the three header tokens plus one move. */
export const MIN_BLOCK = 4;

export const DECLARED_SOURCE = 'la forma declarada en el fichero';
export const RUN_SOURCE = 'la primera respuesta del modelo';

/**
 * The fixed last dimension of the model's only output, or `null` when ORT reports it as symbolic
 * (or reports nothing at all, which is what the WASM backend does for some graphs).
 */
export function declaredVocab(session: SessionLike): number | null {
  const name = session.outputNames[0];
  const meta = session.outputMetadata?.find((entry) => entry.name === name);
  const shape = meta?.shape;
  if (!shape || shape.length === 0) return null;
  const last = shape[shape.length - 1];
  return typeof last === 'number' && Number.isInteger(last) && last > 0 ? last : null;
}

/** Throws with a message a player can read when the logits are not the tokenizer's width. */
export function assertVocab(width: number, expected: number, source: string): void {
  if (width === expected) return;
  throw new Error(
    `El modelo no encaja con el tokenizador: ${source} da ${width} valores y el vocabulario UCI ` +
      `tiene ${expected}. No se juega con este fichero; borra los modelos descargados y vuelve a ` +
      `intentarlo.`,
  );
}

/**
 * The contract of a freshly created session: `block` from the registry (nothing in the session
 * knows it) and `vocab` from the declared output shape when ORT gives a fixed one, checked against
 * the tokenizer right away. When the shape is symbolic the width is still unknown here and
 * `assertVocab` runs on the first `run` instead — see `decoder.worker.ts`.
 */
export function readContract(
  session: SessionLike,
  expectedVocab: number,
  block: number,
): ModelContract {
  if (!Number.isInteger(block) || block < MIN_BLOCK) {
    throw new Error(`El contexto declarado para esta etapa (${block}) no es utilizable.`);
  }
  const declared = declaredVocab(session);
  if (declared !== null) assertVocab(declared, expectedVocab, DECLARED_SOURCE);
  return { block, vocab: declared ?? expectedVocab };
}

// Downloading the model bytes, with the progress bar actually moving while they arrive.
//
// The obvious version — `await cache.put(url, response.clone())` before reading the body — does
// not work: `put` only resolves once the clone's body has been consumed, and a cloned body is fed
// by the same underlying stream, so the browser buffers the whole file (twice: the clone and the
// copy the cache keeps) before the first progress message is ever sent. On an 80 MB fp16 export
// that is a minute of a bar sitting at zero. So the body is streamed first and the cache entry is
// written from the bytes afterwards, which also means a download that fails half-way leaves
// nothing behind.
//
// Everything the browser provides is injected, so the whole thing is testable without a browser.

/** What `downloadModel` needs from its environment. */
export interface DownloadDeps {
  fetch: typeof globalThis.fetch;
  /** `caches`, or undefined where the Cache API is not available (the model is just refetched). */
  caches?: CacheStorage;
  /** Name of the Cache API bucket the weights are stored in. */
  cacheName: string;
  /** Called after every chunk; `total` is the best estimate available at that moment. */
  onProgress: (loaded: number, total: number) => void;
}

/** Joins the chunks a stream produced into one buffer. */
function concat(chunks: readonly Uint8Array[], length: number): Uint8Array {
  const bytes = new Uint8Array(length);
  let offset = 0;
  for (const chunk of chunks) {
    bytes.set(chunk, offset);
    offset += chunk.byteLength;
  }
  return bytes;
}

/**
 * The model bytes, from the Cache API when they are already there and from the network otherwise,
 * reporting progress as the body streams in. `sizeBytes` is the registry's estimate and is only
 * used when the response has no usable `content-length` (a Hub redirect to a CDN sometimes omits
 * it); a `loaded` that overshoots it wins, so the bar never reports more than 100 %.
 */
export async function downloadModel(
  url: string,
  sizeBytes: number,
  deps: DownloadDeps,
): Promise<Uint8Array> {
  const cache = await deps.caches?.open(deps.cacheName).catch(() => null);
  const cached = await cache?.match(url).catch(() => undefined);
  const response = cached ?? (await deps.fetch(url, { mode: 'cors', credentials: 'omit' }));
  if (!response.ok) {
    throw new Error(`no se pudo descargar el modelo (${response.status})`);
  }

  const declared = Number(response.headers.get('content-length') ?? '0');
  const total = declared > 0 ? declared : sizeBytes;

  const body = response.body;
  if (!body) {
    // No stream (an old browser, or a fake in a test): one shot, one progress message.
    const buffer = new Uint8Array(await response.arrayBuffer());
    deps.onProgress(buffer.byteLength, buffer.byteLength);
    return buffer;
  }

  const reader = body.getReader();
  const chunks: Uint8Array[] = [];
  let loaded = 0;
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    loaded += value.byteLength;
    deps.onProgress(loaded, Math.max(total, loaded));
  }
  const bytes = concat(chunks, loaded);
  // The chunks are now a second copy of the whole model: drop them before anyone awaits again.
  chunks.length = 0;
  deps.onProgress(loaded, loaded);

  // Only now, with the whole file in hand and the bar at 100 %, is the cache written.
  if (!cached && cache) {
    // `bytes.buffer` is the freshly allocated buffer `concat` made, so handing it over is a
    // move, not a copy; `Response` reads it synchronously here.
    await cache
      .put(
        url,
        new Response(bytes.buffer as ArrayBuffer, {
          headers: {
            'content-type': 'application/octet-stream',
            'content-length': String(bytes.byteLength),
          },
        }),
      )
      .catch(() => undefined);
  }
  return bytes;
}

// Dedicated worker that owns the ONNX Runtime session: it downloads the model (streaming the
// progress back), keeps it in the Cache API and answers `logits` requests with the last step's
// logits. The main thread never touches ORT, so a 28 MB WASM runtime and a 40 MB model never
// block the board.
//
// Two serialisation rules come from ORT itself and are not negotiable: sessions are created one
// at a time and `run` calls never overlap (the JSEP/asyncify build cannot re-enter an async
// call). Both go through `serial`, a single promise chain — the download included, so two `init`
// messages in flight cannot end up fetching twice and racing to install their session.
//
// The file is never trusted: `readContract` checks the declared output width against the
// tokenizer's vocabulary as soon as the session exists, and `assertVocab` checks the real width
// of every answer. See `src/lib/contract.ts` for why.

import * as ort from 'onnxruntime-web/webgpu';
import { RUN_SOURCE, assertVocab, readContract, type ModelContract } from '../lib/contract';
import { downloadModel } from '../lib/download';
import {
  MODEL_CACHE,
  ORT_BASE,
  type Backend,
  type WorkerRequest,
  type WorkerResponse,
} from '../lib/worker-protocol';

/** The worker global, typed with just what this file uses (avoids pulling in the webworker lib). */
interface WorkerScope {
  postMessage(message: WorkerResponse, transfer?: Transferable[]): void;
  addEventListener(type: 'message', listener: (event: MessageEvent) => void): void;
}

const ctx = self as unknown as WorkerScope;

// Self-hosted runtime: `public/ort/<version>/` is filled by `scripts/copy-assets.mjs` so the
// page never reaches a CDN (the CSP would not allow it either).
ort.env.wasm.wasmPaths = ORT_BASE;
// One thread, always. WASM is only the fallback here (WebGPU is the fast path and this decoder
// is 40 MB), and asking for more threads under cross-origin isolation makes ORT reach for the
// threaded/proxy artefacts that `scripts/copy-assets.mjs` deliberately does not publish. Keeping
// it at one means the runtime needs exactly the two files we do publish, isolated or not.
ort.env.wasm.numThreads = 1;
// No proxy worker either: it would be a third artefact (`ort-wasm-proxy-worker`) and this code
// already runs off the main thread, which is the only thing the proxy buys.
ort.env.wasm.proxy = false;
ort.env.logLevel = 'error';

let session: ort.InferenceSession | null = null;
/** Block and vocabulary the live session was accepted under; null while there is no session. */
let contract: ModelContract | null = null;

/** The single chain every ORT call is queued on: no two sessions or runs are ever in flight. */
let chain: Promise<unknown> = Promise.resolve();

function serial<T>(work: () => Promise<T>): Promise<T> {
  const next = chain.then(work, work);
  chain = next.then(
    () => undefined,
    () => undefined,
  );
  return next;
}

function reply(message: WorkerResponse, transfer?: Transferable[]): void {
  ctx.postMessage(message, transfer);
}

/** `downloadModel` wired to this worker's globals, reporting progress back to the main thread. */
function fetchModel(id: number, url: string, sizeBytes: number): Promise<Uint8Array> {
  return downloadModel(url, sizeBytes, {
    fetch: globalThis.fetch.bind(globalThis),
    caches: typeof caches === 'undefined' ? undefined : caches,
    cacheName: MODEL_CACHE,
    onProgress: (loaded, total) => reply({ type: 'progress', id, loaded, total }),
  });
}

/** The WebGPU adapter, or the reason there is none. Never throws. */
async function webgpuAdapter(): Promise<{ adapter: unknown } | { reason: string }> {
  const gpu = (globalThis as { navigator?: { gpu?: { requestAdapter(): Promise<unknown> } } })
    .navigator?.gpu;
  if (!gpu) return { reason: 'este navegador no expone WebGPU' };
  try {
    const adapter = await gpu.requestAdapter();
    if (!adapter) return { reason: 'WebGPU no ha ofrecido ningún adaptador' };
    return { adapter };
  } catch (cause) {
    return { reason: `WebGPU ha fallado al pedir el adaptador: ${describe(cause)}` };
  }
}

function describe(cause: unknown): string {
  return cause instanceof Error ? cause.message : String(cause);
}

/**
 * WebGPU whenever `requestAdapter()` hands out an adapter; WASM only when it does not, or when
 * creating the session on WebGPU throws. Both fallbacks carry the reason, which the panel shows
 * under the badge so a slow game on WASM is never a mystery.
 */
async function createSession(
  bytes: Uint8Array,
): Promise<{ session: ort.InferenceSession; backend: Backend; reason: string | null }> {
  const options: ort.InferenceSession.SessionOptions = {
    graphOptimizationLevel: 'all',
    executionMode: 'sequential',
  };
  const probe = await webgpuAdapter();
  let reason = 'reason' in probe ? probe.reason : null;
  if (!reason) {
    try {
      const created = await ort.InferenceSession.create(bytes, {
        ...options,
        executionProviders: ['webgpu'],
      });
      return { session: created, backend: 'webgpu', reason: null };
    } catch (cause) {
      // An adapter that cannot compile the graph is a normal outcome, not a bug: say so and
      // carry on with WASM.
      reason = `WebGPU no ha podido crear la sesión: ${describe(cause)}`;
    }
  }
  const created = await ort.InferenceSession.create(bytes, {
    ...options,
    executionProviders: ['wasm'],
  });
  return { session: created, backend: 'wasm', reason };
}

/**
 * Downloads the weights and hands them to ORT. The buffer is a local of this frame on purpose:
 * when it returns, the only copy of the model still alive is the one inside the session, instead
 * of 80 MB sitting next to it for as long as the worker lives.
 */
async function loadSession(request: Extract<WorkerRequest, { type: 'init' }>) {
  const bytes = await fetchModel(request.id, request.url, request.sizeBytes);
  return createSession(bytes);
}

async function init(request: Extract<WorkerRequest, { type: 'init' }>): Promise<void> {
  const started = performance.now();
  // The download runs inside the chain too: outside it, a second `init` would start its own fetch
  // while the first was still tearing the old session down, and both would race to install one.
  const ready = await serial(async () => {
    await session?.release();
    session = null;
    contract = null;
    const created = await loadSession(request);
    const checked = readContract(created.session, request.vocab, request.block);
    session = created.session;
    contract = checked;
    return { backend: created.backend, reason: created.reason, contract: checked };
  });
  reply({
    type: 'ready',
    id: request.id,
    backend: ready.backend,
    fallbackReason: ready.reason ?? undefined,
    loadMs: Math.round(performance.now() - started),
    block: ready.contract.block,
    vocab: ready.contract.vocab,
  });
}

async function logits(request: Extract<WorkerRequest, { type: 'logits' }>): Promise<void> {
  const data = await serial(async () => {
    const current = session;
    const live = contract;
    if (!current || !live) throw new Error('el modelo todavía no está cargado');
    const ids = BigInt64Array.from(request.ids, (id) => BigInt(id));
    const feeds: Record<string, ort.Tensor> = {
      [current.inputNames[0]]: new ort.Tensor('int64', ids, [1, request.ids.length]),
    };
    const started = performance.now();
    const output = await current.run(feeds);
    const inferMs = performance.now() - started;
    const tensor = output[current.outputNames[0]];
    const raw = tensor.data;
    if (!(raw instanceof Float32Array)) {
      throw new Error(`la salida del modelo no es float32 (${tensor.type})`);
    }
    // The declared shape is often symbolic, so the real width is only knowable here. A model
    // whose logits are not the tokenizer's vocabulary would map every id to a different move.
    assertVocab(raw.length, live.vocab, RUN_SOURCE);
    // Copy out of the ORT arena: the tensor's buffer may be reused by the next run.
    return { values: raw.slice(), inferMs };
  });
  reply({ type: 'logits', id: request.id, data: data.values, inferMs: data.inferMs }, [
    data.values.buffer,
  ]);
}

async function dispose(request: Extract<WorkerRequest, { type: 'dispose' }>): Promise<void> {
  await serial(async () => {
    await session?.release();
    session = null;
    contract = null;
  });
  reply({ type: 'disposed', id: request.id });
}

ctx.addEventListener('message', (event: MessageEvent) => {
  const request = event.data as WorkerRequest;
  const run = async () => {
    switch (request.type) {
      case 'init':
        return init(request);
      case 'logits':
        return logits(request);
      case 'dispose':
        return dispose(request);
      default:
        return undefined;
    }
  };
  void run().catch((error: unknown) => {
    reply({
      type: 'error',
      id: request.id,
      message: error instanceof Error ? error.message : String(error),
    });
  });
});

// Main-thread side of the decoder worker. The `new Worker(new URL(...), { type: 'module' })`
// literal lives here, at the call site, because that is the only shape Vite recognises to emit
// the worker as its own chunk (`tests/worker-chunk.test.ts` guards it).
import {
  asResponse,
  type Backend,
  type InitRequest,
  type ProgressMessage,
  type ReadyMessage,
  type WorkerRequest,
} from '../lib/worker-protocol';

export interface LoadReport {
  backend: Backend;
  /** Why WebGPU was not used, when it was not. */
  fallbackReason?: string;
  loadMs: number;
  /** Context window the session was accepted under; `buildPrompt` crops the prompt to it. */
  block: number;
  /** Width of the logits the model answers with, already checked against the tokenizer. */
  vocab: number;
}

/** What `init` needs to know about a stage; the worker's `id` is added by `send`. */
export type LoadRequest = Omit<InitRequest, 'type' | 'id'>;

export interface LogitsReport {
  data: Float32Array;
  inferMs: number;
}

export interface Decoder {
  /** Downloads (or reads from the cache) the model, creates the session and checks its contract. */
  init(request: LoadRequest, onProgress?: (progress: ProgressMessage) => void): Promise<LoadReport>;
  /** Logits of the last step for a prompt already cropped to the context window. */
  logits(ids: number[]): Promise<LogitsReport>;
  /** Releases the session and terminates the worker; the handle is unusable afterwards. */
  dispose(): Promise<void>;
}

/** `Omit` over a union collapses it, so the distribution is written out explicitly. */
type Unidentified<T> = T extends WorkerRequest ? Omit<T, 'id'> : never;

interface Pending {
  resolve: (value: never) => void;
  reject: (error: Error) => void;
  onProgress?: (progress: ProgressMessage) => void;
}

export function createDecoder(): Decoder {
  const worker = new Worker(new URL('./decoder.worker.ts', import.meta.url), { type: 'module' });
  const pending = new Map<number, Pending>();
  let nextId = 1;
  let terminated = false;

  const fail = (error: Error) => {
    for (const entry of pending.values()) entry.reject(error);
    pending.clear();
  };

  worker.addEventListener('message', (event: MessageEvent) => {
    const message = asResponse(event.data);
    if (!message) return;
    const entry = pending.get(message.id);
    if (!entry) return;
    if (message.type === 'progress') {
      entry.onProgress?.(message);
      return;
    }
    pending.delete(message.id);
    if (message.type === 'error') {
      entry.reject(new Error(message.message));
      return;
    }
    (entry.resolve as (value: unknown) => void)(message);
  });

  worker.addEventListener('error', (event: ErrorEvent) => {
    fail(new Error(event.message || 'el worker del modelo ha fallado'));
  });

  function send<T>(
    request: Unidentified<WorkerRequest>,
    onProgress?: (progress: ProgressMessage) => void,
  ): Promise<T> {
    if (terminated) return Promise.reject(new Error('el worker del modelo ya se ha cerrado'));
    const id = nextId++;
    return new Promise<T>((resolve, reject) => {
      pending.set(id, { resolve: resolve as (value: never) => void, reject, onProgress });
      worker.postMessage({ ...request, id } as WorkerRequest);
    });
  }

  return {
    async init(request, onProgress) {
      const ready = await send<ReadyMessage>({ type: 'init', ...request }, onProgress);
      return {
        backend: ready.backend,
        fallbackReason: ready.fallbackReason,
        loadMs: ready.loadMs,
        block: ready.block,
        vocab: ready.vocab,
      };
    },
    async logits(ids) {
      const answer = await send<{ data: Float32Array; inferMs: number }>({ type: 'logits', ids });
      return { data: answer.data, inferMs: answer.inferMs };
    },
    async dispose() {
      if (terminated) return;
      try {
        await send<unknown>({ type: 'dispose' });
      } catch {
        /* the worker may already be gone; terminating below is enough */
      }
      terminated = true;
      fail(new Error('el worker del modelo se ha cerrado'));
      worker.terminate();
    },
  };
}

export type StageKind = 'mock' | 'onnx';

export interface Stage {
  id: string;
  label: string;
  kind: StageKind;
  /** Download size in MB, shown in the consent step. 0 for stages that download nothing. */
  sizeMb: number;
  /** Hugging Face repo id, for `kind: 'onnx'` stages served from the Hub. */
  repo?: string;
  /** Path of the ONNX file inside the repo. */
  file?: string;
  /** Same-origin URL that replaces `repo`/`file` (only the toy model used by the E2E suite). */
  url?: string;
  /** True once the stage is trained with the Elo conditioning tokens (M4). */
  eloConditioned?: boolean;
  /**
   * Context window the stage was trained with (`DecoderConfig.block`). It lives here because ORT
   * Web does not expose the `rukh_block` metadata the exporter writes into the file, so the only
   * place the browser can learn it is the registry; `DEFAULT_BLOCK` when a stage omits it.
   */
  block?: number;
}

/** `DecoderConfig.block` of every model published so far. */
export const DEFAULT_BLOCK = 200;

/** The context window a stage plays with. */
export function stageBlock(stage: Stage): number {
  return stage.block ?? DEFAULT_BLOCK;
}

/**
 * Download sizes in MB. **Provisional**: they are the plan's estimates for the fp16 and int8
 * exports and the controller updates them here, in this one place, after the real export
 * (`rukh export ... --fp16 --int8`) reports the file sizes. They only drive the consent copy and
 * the progress bar fallback, never the download itself.
 */
export const STAGE_SIZE_MB = {
  'tiny-int8': 6,
  'small-fp16': 80,
  'small-int8': 40,
} as const;

/** Licence of every published Rukh model; shown before anything is downloaded. */
export const MODEL_LICENSE = 'Apache-2.0';

export const STAGES: Stage[] = [
  { id: 'mock', label: 'Primera jugada legal', kind: 'mock', sizeMb: 0 },
  {
    id: 'tiny-int8',
    label: 'Rukh tiny (int8)',
    kind: 'onnx',
    repo: 'chorcat/rukh-tiny',
    file: 'onnx/model-int8.onnx',
    sizeMb: STAGE_SIZE_MB['tiny-int8'],
    block: DEFAULT_BLOCK,
  },
  {
    id: 'small-fp16',
    label: 'Rukh small (fp16)',
    kind: 'onnx',
    repo: 'chorcat/rukh-small',
    file: 'onnx/model-fp16.onnx',
    sizeMb: STAGE_SIZE_MB['small-fp16'],
    block: DEFAULT_BLOCK,
  },
  {
    id: 'small-int8',
    label: 'Rukh small (int8)',
    kind: 'onnx',
    repo: 'chorcat/rukh-small',
    file: 'onnx/model-int8.onnx',
    sizeMb: STAGE_SIZE_MB['small-int8'],
    block: DEFAULT_BLOCK,
  },
];

/**
 * Toy decoder committed under `public/test/`: one layer, `d_model=8`, and the real contract
 * (`idx (B, T)` int64 -> `logits (B, 2030)` with `block` 200), exported by the very function that
 * writes the published models (`rukh.export.export_onnx`, dynamo exporter, `rukh_*` metadata). It
 * is not in `STAGES` (nothing offers it in the selector); only `?stage=test` reaches it, which is
 * how the E2E suite exercises the real worker path — contract check included — without
 * downloading 40 MB from the Hub.
 */
export const TEST_STAGE: Stage = {
  id: 'test',
  label: 'ONNX de juguete (pruebas)',
  kind: 'onnx',
  sizeMb: 0.2,
  url: '/test/toy-decoder.onnx',
  block: DEFAULT_BLOCK,
};

export function findStage(id: string): Stage | undefined {
  if (id === TEST_STAGE.id) return TEST_STAGE;
  return STAGES.find((stage) => stage.id === id);
}

/** Hub URL of a stage, or its own `url` when it is served from this origin. */
export function modelUrl(stage: Stage): string {
  if (stage.url) return stage.url;
  if (!stage.repo || !stage.file) {
    throw new Error(`stage ${stage.id} has no model file`);
  }
  return `https://huggingface.co/${stage.repo}/resolve/main/${stage.file}`;
}

/** The stages the selector offers: the registry plus `current` when it is not part of it. */
export function selectableStages(current: string): Stage[] {
  const extra = STAGES.some((stage) => stage.id === current) ? undefined : findStage(current);
  return extra ? [...STAGES, extra] : STAGES;
}

export interface Conditions {
  /** `navigator.connection?.saveData`. */
  saveData?: boolean;
  /** True on a phone-sized or coarse-pointer device. */
  mobile?: boolean;
}

/**
 * Stage selected when the page is opened without `?stage=`: the 40 MB int8 export on mobile or
 * with data saver on, the 80 MB fp16 one elsewhere.
 */
export function defaultStageId(conditions: Conditions = {}): string {
  return conditions.saveData || conditions.mobile ? 'small-int8' : 'small-fp16';
}

/** `defaultStageId` reading the browser it runs in; falls back to the desktop choice. */
export function detectConditions(): Conditions {
  if (typeof navigator === 'undefined') return {};
  const connection = (navigator as Navigator & { connection?: { saveData?: boolean } }).connection;
  const mobile =
    typeof window !== 'undefined' && typeof window.matchMedia === 'function'
      ? window.matchMedia('(max-width: 820px), (pointer: coarse)').matches
      : false;
  return { saveData: connection?.saveData === true, mobile };
}

# Included in every location: nginx does not inherit add_header into a block that declares
# its own add_header, so each location pulls the full set from here.
#
# The page-level CSP (script/style hashes) lives in the <meta> tag generated by Astro;
# frame-ancestors cannot be set from <meta>, so it is added as a header.
add_header Content-Security-Policy "frame-ancestors 'none'" always;
add_header Strict-Transport-Security "max-age=31536000; includeSubDomains" always;
add_header X-Content-Type-Options "nosniff" always;
add_header Referrer-Policy "strict-origin-when-cross-origin" always;
add_header Permissions-Policy "camera=(), microphone=(), geolocation=(), payment=(), usb=()" always;
# Cross-origin isolation: required for multi-threaded WASM (onnxruntime-web, P2).
add_header Cross-Origin-Opener-Policy "same-origin" always;
add_header Cross-Origin-Embedder-Policy "require-corp" always;