src/engine/native-log.ts

export interface GpuLayerReport {
  offloaded: number;
  total: number;
}

export interface BackendReport {
  backends: string[];
  nGraphSplits: number | null;
  opsOnCpu: number;
  layersGpu: GpuLayerReport | null;
  flashAttention: boolean | null;
  cacheTypeK: string | null;
  cacheTypeV: string | null;
  webgpuKvBufferBytes: number | null;
  modelBufferBytes?: number;
  computeBufferBytes?: number;
  outputBufferBytes?: number;
  allocatedBufferBytes?: number;
  webgpuTrace: string[];
}

export interface NativeDeviceLostSignal {
  line: string;
  reason: number | null;
  message: string | null;
}

export interface NativeDeviceLostWatch {
  promise: Promise<NativeDeviceLostSignal>;
  dispose(): void;
}

export interface NativeModelLoadProgress {
  state: 'loading' | 'ready';
  stages: string[];
  current: string | null;
  value: number;
}

const BACKEND_NAMES = ['WebGPU', 'CPU', 'Metal', 'CUDA', 'Vulkan'] as const;
const EVIDENCE = /using device|model buffer size|KV buffer size|compute buffer size|output buffer size|offloaded \d+\/\d+ layers|graph splits\s*=|graph_compute\(|backend\s*=\s*CPU|assigned to (?:device )?CPU|(?:ops?|nodes?).*(?:on|to) CPU|flash_attn\s*=|flash attention.*(?:enabled|disabled)|\bK\s*\([^)]+\).*\bV\s*\([^)]+\)/i;
const WEBGPU_DEVICE_LOST = /\bdevice\s+lost!?(?:\s+reason:\s*(-?\d+),\s*message:\s*(.*))?/i;
const WEBGPU_KV_BUFFER = /\bWebGPU\s+KV buffer size\s*=\s*(\d+(?:\.\d+)?)\s*(KiB|MiB|GiB)\b/i;
const WEBGPU_TRACE_PREFIX = '@@WEBGPU_TRACE@@';
const MODEL_LOAD_PROGRESS_PREFIX = '@@MODEL_LOAD_PROGRESS@@';
const WEBGPU_TRACE_LINE_LIMIT = 16_384;

const BINARY_UNIT_BYTES = {
  kib: 1024,
  mib: 1024 ** 2,
  gib: 1024 ** 3,
} as const;

function canonicalBackend(value: string): string {
  const match = BACKEND_NAMES.find((backend) => value.toLowerCase().startsWith(backend.toLowerCase()));
  return match ?? value;
}

function cpuOpsInLine(line: string): number {
  const nodePatterns = [
    /ggml_backend_cpu_graph_compute\(\s*(\d+)\s+nodes?\s*\)/i,
    /backend\s*=\s*CPU\b.*?n_(?:nodes|ops)\s*=\s*(\d+)/i,
    /(?:backend|device)\s+CPU\b.*?(?:graph|compute).*?(\d+)\s+(?:nodes?|ops?)/i,
    /(\d+)\s+(?:nodes?|ops?).*?(?:on|to)\s+CPU\b/i,
  ];
  for (const pattern of nodePatterns) {
    const match = line.match(pattern);
    if (match) {
      return Number(match[1]);
    }
  }
  if (/\b(?:op|node)\b.*?assigned to (?:device )?CPU\b/i.test(line)) {
    return 1;
  }
  return 0;
}

export function isNativeLogEvidenceLine(line: string): boolean {
  return EVIDENCE.test(line);
}

export function parseNativeDeviceLostSignal(line: string): NativeDeviceLostSignal | null {
  if (!/\b(?:ggml[_ -]?webgpu|webgpu)\b/i.test(line)) {
    return null;
  }
  const match = line.match(WEBGPU_DEVICE_LOST);
  if (!match) {
    return null;
  }
  const reason = match[1] === undefined ? null : Number(match[1]);
  const message = match[2]?.trim() || null;
  return {
    line: line.slice(0, 4096),
    reason: Number.isSafeInteger(reason) ? reason : null,
    message,
  };
}

export function parseNativeModelLoadProgress(line: string): NativeModelLoadProgress | null {
  if (!line.startsWith(MODEL_LOAD_PROGRESS_PREFIX)) return null;
  try {
    const payload = JSON.parse(line.slice(MODEL_LOAD_PROGRESS_PREFIX.length)) as Record<string, unknown>;
    if (payload.state !== 'loading' && payload.state !== 'ready') return null;
    const stages = Array.isArray(payload.stages)
      ? payload.stages.filter((stage): stage is string => typeof stage === 'string')
      : [];
    const current = typeof payload.current === 'string' ? payload.current : null;
    const value = payload.state === 'ready'
      ? 1
      : typeof payload.value === 'number' && Number.isFinite(payload.value)
        ? Math.min(1, Math.max(0, payload.value))
        : 0;
    return { state: payload.state, stages, current, value };
  } catch {
    return null;
  }
}

export function parseNativeLog(input: string | readonly string[]): BackendReport {
  const lines = typeof input === 'string' ? input.split(/\r?\n/) : input;
  const backends = new Set<string>();
  const graphSplits: number[] = [];
  let opsOnCpu = 0;
  let layersGpu: GpuLayerReport | null = null;
  let flashAttention: boolean | null = null;
  let cacheTypeK: string | null = null;
  let cacheTypeV: string | null = null;
  let webgpuKvBufferBytes: number | null = null;
  let modelBufferBytes = 0;
  let computeBufferBytes = 0;
  let outputBufferBytes = 0;

  for (const line of lines) {
    const deviceMatch = line.match(/using device\s+([^\s(]+)\s+\(([^)]+)\)/i);
    if (deviceMatch) {
      backends.add(canonicalBackend(deviceMatch[2] ?? deviceMatch[1] ?? 'unknown'));
    }

    for (const backend of BACKEND_NAMES) {
      const escaped = backend.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
      if (new RegExp(`\\b${escaped}\\s+(?:model|KV|compute|output) buffer size`, 'i').test(line)) {
        backends.add(backend);
      }
    }

    const layerMatch = line.match(/offloaded\s+(\d+)\s*\/\s*(\d+)\s+layers to GPU/i);
    if (layerMatch) {
      layersGpu = {
        offloaded: Number(layerMatch[1]),
        total: Number(layerMatch[2]),
      };
    }

    const graphMatch = line.match(/graph splits\s*=\s*(\d+)/i);
    if (graphMatch) {
      graphSplits.push(Number(graphMatch[1]));
    }

    const contextFlashMatch = line.match(/\bflash_attn\s*=\s*(enabled|disabled)\b/i);
    const resolvedFlashMatch = line.match(/\bFlash Attention\b.*\b(enabled|disabled)\b/i);
    const flashState = (contextFlashMatch ?? resolvedFlashMatch)?.[1];
    if (flashState) {
      flashAttention = flashState.toLowerCase() === 'enabled';
    }

    const cacheTypesMatch = line.match(/\bK\s*\(\s*([^)]+?)\s*\).*\bV\s*\(\s*([^)]+?)\s*\)/i);
    const cacheK = cacheTypesMatch?.[1];
    const cacheV = cacheTypesMatch?.[2];
    if (cacheK && cacheV) {
      cacheTypeK = cacheK.trim().toLowerCase();
      cacheTypeV = cacheV.trim().toLowerCase();
    }

    const kvBufferMatch = line.match(WEBGPU_KV_BUFFER);
    const kvBufferSize = kvBufferMatch?.[1];
    const kvBufferUnit = kvBufferMatch?.[2];
    if (kvBufferSize && kvBufferUnit) {
      const unitBytes = BINARY_UNIT_BYTES[kvBufferUnit.toLowerCase() as keyof typeof BINARY_UNIT_BYTES];
      const bytes = Math.round(Number(kvBufferSize) * unitBytes);
      webgpuKvBufferBytes = (webgpuKvBufferBytes ?? 0) + bytes;
    }
    const bufferMatch = line.match(/\b(model|compute|output) buffer size\s*=\s*(\d+(?:\.\d+)?)\s*(KiB|MiB|GiB)\b/i);
    const bufferKind = bufferMatch?.[1]?.toLowerCase();
    const bufferSize = bufferMatch?.[2];
    const bufferUnit = bufferMatch?.[3];
    if (bufferKind && bufferSize && bufferUnit) {
      const unitBytes = BINARY_UNIT_BYTES[bufferUnit.toLowerCase() as keyof typeof BINARY_UNIT_BYTES];
      const bytes = Math.round(Number(bufferSize) * unitBytes);
      if (bufferKind === 'model') modelBufferBytes += bytes;
      if (bufferKind === 'compute') computeBufferBytes += bytes;
      if (bufferKind === 'output') outputBufferBytes += bytes;
    }
    opsOnCpu += cpuOpsInLine(line);
  }

  const allocatedBufferBytes = modelBufferBytes
    + computeBufferBytes
    + outputBufferBytes
    + (webgpuKvBufferBytes ?? 0);
  return {
    backends: [...backends],
    nGraphSplits: graphSplits.length > 0 ? Math.max(...graphSplits) : null,
    opsOnCpu,
    layersGpu,
    flashAttention,
    cacheTypeK,
    cacheTypeV,
    webgpuKvBufferBytes,
    ...(modelBufferBytes > 0 ? { modelBufferBytes } : {}),
    ...(computeBufferBytes > 0 ? { computeBufferBytes } : {}),
    ...(outputBufferBytes > 0 ? { outputBufferBytes } : {}),
    ...(allocatedBufferBytes > 0 ? { allocatedBufferBytes } : {}),
    webgpuTrace: [],
  };
}

export class NativeLogCollector {
  private lines: string[] = [];
  private webgpuTrace: string[] = [];
  private recentLines: string[] = [];
  private deviceLost: NativeDeviceLostSignal | null = null;
  private readonly deviceLostListeners = new Set<(signal: NativeDeviceLostSignal) => void>();
  private readonly modelLoadListeners = new Set<(progress: NativeModelLoadProgress) => void>();

  clear(): void {
    this.lines = [];
    this.webgpuTrace = [];
    this.recentLines = [];
    this.deviceLost = null;
  }

  append(values: readonly unknown[]): void {
    const line = values
      .map((value) => (typeof value === 'string' ? value : JSON.stringify(value)))
      .join(' ');
    if (line) {
      this.recentLines.push(line.slice(0, 4096));
      if (this.recentLines.length > 64) {
        this.recentLines.shift();
      }
    }
    const modelLoadProgress = parseNativeModelLoadProgress(line);
    if (modelLoadProgress) {
      for (const listener of this.modelLoadListeners) listener(modelLoadProgress);
    }
    if (!this.deviceLost) {
      const deviceLost = parseNativeDeviceLostSignal(line);
      if (deviceLost) {
        this.deviceLost = deviceLost;
        for (const listener of this.deviceLostListeners) listener(deviceLost);
        this.deviceLostListeners.clear();
      }
    }
    if (isNativeLogEvidenceLine(line)) {
      this.lines.push(line.slice(0, 4096));
      if (this.lines.length > 512) {
        this.lines.shift();
      }
    }
    if (line.startsWith(WEBGPU_TRACE_PREFIX)) {
      this.webgpuTrace.push(line.slice(0, 4096));
      if (this.webgpuTrace.length > WEBGPU_TRACE_LINE_LIMIT) {
        this.webgpuTrace.shift();
      }
    }
  }

  report(): BackendReport {
    return {
      ...parseNativeLog(this.lines),
      webgpuTrace: [...this.webgpuTrace],
    };
  }

  recent(limit = 16): string[] {
    return this.recentLines.slice(-Math.max(0, limit));
  }

  deviceLostSignal(): NativeDeviceLostSignal | null {
    return this.deviceLost;
  }

  watchDeviceLost(): NativeDeviceLostWatch {
    if (this.deviceLost) {
      return { promise: Promise.resolve(this.deviceLost), dispose: () => undefined };
    }
    let listener: ((signal: NativeDeviceLostSignal) => void) | null = null;
    const promise = new Promise<NativeDeviceLostSignal>((resolve) => {
      listener = resolve;
      this.deviceLostListeners.add(resolve);
    });
    return {
      promise,
      dispose: () => {
        if (listener) this.deviceLostListeners.delete(listener);
        listener = null;
      },
    };
  }

  onModelLoadProgress(listener: (progress: NativeModelLoadProgress) => void): () => void {
    this.modelLoadListeners.add(listener);
    return () => this.modelLoadListeners.delete(listener);
  }
}