src/engine/manifest.ts

export type ModelTierId = '1_7b' | '4b' | '8b' | '27b';

export interface ManifestSourceModel {
  repo: string;
  revision: string;
  file: string;
  bytes: number;
  sha256: string;
}

export interface ManifestModelFile {
  path: string;
  bytes: number;
  sha256: string;
  url?: string;
}

export interface ManifestChatTemplate {
  bytes: number;
  sha256: string;
  markers: {
    think: boolean;
    toolCall: boolean;
    toolResponse: boolean;
  };
}

export interface ManifestModelV2 {
  id: ModelTierId;
  displayName: string;
  architecture: string;
  source: ManifestSourceModel;
  files: ManifestModelFile[];
  downloadBytes: number;
  contextLength: number;
  defaultContext: number;
  cpuFallback: boolean;
  largestTensorBytes: number;
  requiredLimits: Record<string, number>;
  chatTemplate: ManifestChatTemplate;
  hybridDimensions: Record<string, number>;
  nextNTensorCount: number;
  runtimePolicy: {
    flashAttention: false;
    tokenEmbeddingOnWebGPU: boolean;
    requireSingleWebGPUGraph: boolean;
  };
}

export interface ModelManifestV2 {
  schemaVersion: 2;
  repository: {
    id: string;
    revision: string | null;
  };
  models: ManifestModelV2[];
}

const MODEL_IDS = new Set<ModelTierId>(['1_7b', '4b', '8b', '27b']);
const SHA256 = /^[0-9a-f]{64}$/;
const REVISION = /^[0-9a-f]{40}$/;
const SPLIT_FILE = /-(\d{5})-of-(\d{5})\.gguf$/i;

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === 'object' && value !== null && !Array.isArray(value);
}

function assert(condition: unknown, message: string): asserts condition {
  if (!condition) {
    throw new Error(`Invalid model manifest v2: ${message}`);
  }
}

function requireString(value: unknown, field: string): string {
  assert(typeof value === 'string' && value.length > 0, `${field} must be a non-empty string`);
  return value;
}

function requirePositiveInteger(value: unknown, field: string): number {
  assert(Number.isSafeInteger(value) && Number(value) > 0, `${field} must be a positive integer`);
  return Number(value);
}

function requireNonNegativeInteger(value: unknown, field: string): number {
  assert(Number.isSafeInteger(value) && Number(value) >= 0, `${field} must be a non-negative integer`);
  return Number(value);
}

function requireSha256(value: unknown, field: string): string {
  const sha256 = requireString(value, field);
  assert(SHA256.test(sha256), `${field} must be lowercase SHA-256 hex`);
  return sha256;
}

function requireRevision(value: unknown, field: string): string {
  const revision = requireString(value, field);
  assert(REVISION.test(revision), `${field} must be a full 40-character commit`);
  return revision;
}

function requireRepo(value: unknown, field: string): string {
  const repo = requireString(value, field);
  assert(/^[^/\s]+\/[^/\s]+$/.test(repo), `${field} must be a Hugging Face owner/repository id`);
  return repo;
}

function requirePath(value: unknown, field: string): string {
  const path = requireString(value, field);
  assert(!path.startsWith('/'), `${field} must be repository-relative`);
  assert(!path.split('/').some((part) => part === '' || part === '.' || part === '..'), `${field} contains an unsafe segment`);
  assert(path.toLowerCase().endsWith('.gguf'), `${field} must point to a GGUF file`);
  return path;
}

function parseSource(value: unknown, field: string): ManifestSourceModel {
  assert(isRecord(value), `${field} must be an object`);
  return {
    repo: requireRepo(value.repo, `${field}.repo`),
    revision: requireRevision(value.revision, `${field}.revision`),
    file: requirePath(value.file, `${field}.file`),
    bytes: requirePositiveInteger(value.bytes, `${field}.bytes`),
    sha256: requireSha256(value.sha256, `${field}.sha256`),
  };
}

function parseFile(value: unknown, field: string): ManifestModelFile {
  assert(isRecord(value), `${field} must be an object`);
  const file: ManifestModelFile = {
    path: requirePath(value.path, `${field}.path`),
    bytes: requirePositiveInteger(value.bytes, `${field}.bytes`),
    sha256: requireSha256(value.sha256, `${field}.sha256`),
  };
  if (value.url !== undefined) {
    const url = requireString(value.url, `${field}.url`);
    const parsed = new URL(url);
    assert(parsed.protocol === 'https:', `${field}.url must use HTTPS`);
    file.url = url;
  }
  return file;
}

function parsePositiveNumberRecord(value: unknown, field: string): Record<string, number> {
  assert(isRecord(value), `${field} must be an object`);
  const parsed: Record<string, number> = {};
  for (const [key, raw] of Object.entries(value)) {
    assert(key.length > 0, `${field} contains an empty key`);
    parsed[key] = requirePositiveInteger(raw, `${field}.${key}`);
  }
  return parsed;
}

function parseChatTemplate(value: unknown, field: string): ManifestChatTemplate {
  assert(isRecord(value), `${field} must be an object`);
  assert(isRecord(value.markers), `${field}.markers must be an object`);
  const markers = value.markers;
  assert(typeof markers.think === 'boolean', `${field}.markers.think must be boolean`);
  assert(typeof markers.toolCall === 'boolean', `${field}.markers.toolCall must be boolean`);
  assert(typeof markers.toolResponse === 'boolean', `${field}.markers.toolResponse must be boolean`);
  return {
    bytes: requirePositiveInteger(value.bytes, `${field}.bytes`),
    sha256: requireSha256(value.sha256, `${field}.sha256`),
    markers: {
      think: markers.think,
      toolCall: markers.toolCall,
      toolResponse: markers.toolResponse,
    },
  };
}

function validateFileOrder(files: ManifestModelFile[], field: string): void {
  const matches = files.map((file) => file.path.match(SPLIT_FILE));
  const splitCount = matches.filter(Boolean).length;
  if (splitCount === 0) {
    assert(files.length === 1, `${field} must use GGUF split names when it contains multiple files`);
    return;
  }
  assert(splitCount === files.length, `${field} mixes split and non-split GGUF names`);
  for (let index = 0; index < matches.length; index += 1) {
    const match = matches[index];
    assert(match !== undefined && match !== null, `${field}[${index}] is not a split GGUF name`);
    assert(Number(match[1]) === index + 1, `${field}[${index}] is out of order`);
    assert(Number(match[2]) === files.length, `${field}[${index}] has the wrong shard count`);
  }
}

function parseModel(value: unknown, index: number): ManifestModelV2 {
  const field = `models[${index}]`;
  assert(isRecord(value), `${field} must be an object`);
  const id = requireString(value.id, `${field}.id`);
  assert(MODEL_IDS.has(id as ModelTierId), `${field}.id is not a supported Bonsai tier`);
  assert(Array.isArray(value.files) && value.files.length > 0, `${field}.files must be a non-empty array`);
  const files = value.files.map((file, fileIndex) => parseFile(file, `${field}.files[${fileIndex}]`));
  validateFileOrder(files, `${field}.files`);
  assert(new Set(files.map((file) => file.path)).size === files.length, `${field}.files contains duplicate paths`);

  const downloadBytes = requirePositiveInteger(value.downloadBytes, `${field}.downloadBytes`);
  assert(
    files.reduce((sum, file) => sum + file.bytes, 0) === downloadBytes,
    `${field}.downloadBytes does not equal the shard byte sum`,
  );
  const contextLength = requirePositiveInteger(value.contextLength, `${field}.contextLength`);
  const defaultContext = requirePositiveInteger(value.defaultContext, `${field}.defaultContext`);
  assert(defaultContext <= contextLength, `${field}.defaultContext exceeds contextLength`);
  assert(typeof value.cpuFallback === 'boolean', `${field}.cpuFallback must be boolean`);
  assert(isRecord(value.runtimePolicy), `${field}.runtimePolicy must be an object`);
  assert(value.runtimePolicy.flashAttention === false, `${field}.runtimePolicy.flashAttention must be false`);
  assert(
    typeof value.runtimePolicy.tokenEmbeddingOnWebGPU === 'boolean',
    `${field}.runtimePolicy.tokenEmbeddingOnWebGPU must be boolean`,
  );
  assert(
    typeof value.runtimePolicy.requireSingleWebGPUGraph === 'boolean',
    `${field}.runtimePolicy.requireSingleWebGPUGraph must be boolean`,
  );

  const model: ManifestModelV2 = {
    id: id as ModelTierId,
    displayName: requireString(value.displayName, `${field}.displayName`),
    architecture: requireString(value.architecture, `${field}.architecture`),
    source: parseSource(value.source, `${field}.source`),
    files,
    downloadBytes,
    contextLength,
    defaultContext,
    cpuFallback: value.cpuFallback,
    largestTensorBytes: requirePositiveInteger(value.largestTensorBytes, `${field}.largestTensorBytes`),
    requiredLimits: parsePositiveNumberRecord(value.requiredLimits, `${field}.requiredLimits`),
    chatTemplate: parseChatTemplate(value.chatTemplate, `${field}.chatTemplate`),
    hybridDimensions: parsePositiveNumberRecord(value.hybridDimensions, `${field}.hybridDimensions`),
    nextNTensorCount: requireNonNegativeInteger(value.nextNTensorCount, `${field}.nextNTensorCount`),
    runtimePolicy: {
      flashAttention: false,
      tokenEmbeddingOnWebGPU: value.runtimePolicy.tokenEmbeddingOnWebGPU,
      requireSingleWebGPUGraph: value.runtimePolicy.requireSingleWebGPUGraph,
    },
  };

  assert(
    model.requiredLimits.maxStorageBufferBindingSize === model.largestTensorBytes,
    `${field}.requiredLimits.maxStorageBufferBindingSize must equal largestTensorBytes`,
  );
  if (model.id === '27b') {
    assert(model.cpuFallback === false, `${field} must disable CPU fallback for 27B`);
    assert(model.runtimePolicy.requireSingleWebGPUGraph, `${field} must require a single WebGPU graph for 27B`);
  }
  return model;
}

export function parseModelManifestV2(value: unknown): ModelManifestV2 {
  assert(isRecord(value), 'root must be an object');
  assert(value.schemaVersion === 2, 'schemaVersion must equal 2');
  assert(isRecord(value.repository), 'repository must be an object');
  const repositoryId = requireRepo(value.repository.id, 'repository.id');
  const repositoryRevision = value.repository.revision === null
    ? null
    : requireRevision(value.repository.revision, 'repository.revision');
  assert(Array.isArray(value.models), 'models must be an array');
  const models = value.models.map(parseModel);
  assert(models.length === MODEL_IDS.size, 'models must contain exactly the four Bonsai tiers');
  assert(new Set(models.map((model) => model.id)).size === models.length, 'models contains duplicate ids');
  assert(models.every((model) => MODEL_IDS.has(model.id)), 'models contains an unsupported id');
  return {
    schemaVersion: 2,
    repository: {
      id: repositoryId,
      revision: repositoryRevision,
    },
    models,
  };
}

export async function loadModelManifestV2(url: string, signal?: AbortSignal): Promise<ModelManifestV2> {
  const response = await fetch(url, {
    ...(signal ? { signal } : {}),
    headers: { Accept: 'application/json' },
  });
  if (!response.ok) {
    throw new Error(`Failed to load model manifest: HTTP ${response.status}`);
  }
  return parseModelManifestV2(await response.json());
}

function encodeRepositoryPath(path: string): string {
  return path.split('/').map((part) => encodeURIComponent(part)).join('/');
}

export function orderedShardUrls(manifest: ModelManifestV2, model: ManifestModelV2): string[] {
  assert(
    manifest.repository.revision !== null,
    'repository.revision is required before model shards can be loaded',
  );
  return model.files.map((file) => {
    const derived = `https://huggingface.co/${manifest.repository.id}/resolve/${manifest.repository.revision}/${encodeRepositoryPath(file.path)}`;
    if (file.url) {
      const explicit = new URL(file.url);
      const expected = new URL(derived);
      assert(
        explicit.origin === expected.origin && explicit.pathname === expected.pathname,
        `${file.path} URL must match the commit-pinned repository revision`,
      );
      return file.url;
    }
    return derived;
  });
}

export function findManifestModel(manifest: ModelManifestV2, id: ModelTierId): ManifestModelV2 {
  const model = manifest.models.find((candidate) => candidate.id === id);
  if (!model) {
    throw new Error(`Model ${id} is missing from the manifest`);
  }
  return model;
}