import { MODEL, DSPARK, DEFAULTS } from "./core.js";
import { WeightCache } from "./cache.js";
import { downloadForSession } from "./session-download.js";
import { startRuntime } from "./startup.js";
import { splitStreamContent } from "./stream.js";
export class Engine {
  constructor(onStatus) {
    this.onStatus = onStatus;
    this.logs = [];
    this.ready = false;
    this.loading = false;
    this.adapter = null;
    this.gpuLayers = 0;
    this.cache = new WeightCache();
    this.sessionFiles = new Map();
  }
  async probe() {
    if (!navigator.gpu)
      return {
        ok: false,
        message: "WebGPU is unavailable. Use a browser with WebGPU enabled.",
      };
    this.adapter = await navigator.gpu.requestAdapter({
      powerPreference: "high-performance",
    });
    if (!this.adapter)
      return {
        ok: false,
        message:
          "No WebGPU adapter found. Enable hardware acceleration in your browser, then reopen this page.",
      };
    const info = this.adapter.info;
    return {
      ok: true,
      name:
        info?.description || info?.device || info?.vendor || "WebGPU graphics",
      maxBuffer: this.adapter.limits.maxStorageBufferBindingSize,
      features: [...this.adapter.features],
      mobile: /Android|iPhone|iPad/i.test(navigator.userAgent),
    };
  }
  get url() {
    if (
      import.meta.env.DEV &&
      new URLSearchParams(location.search).get("qa-weights") === "local"
    )
      return location.origin + "/__qa/model/" + MODEL.file;
    return `https://huggingface.co/${MODEL.weights}/resolve/${MODEL.revision}/${MODEL.file}`;
  }
  get urls() {
    return Array.from({ length: 6 }, (_, i) =>
      this.url.replace("-00001-of-", `-${String(i + 1).padStart(5, "0")}-of-`),
    );
  }
  get draftURL() {
    return import.meta.env.DEV &&
      new URLSearchParams(location.search).get("qa-weights") === "local"
      ? location.origin + "/__qa/dspark"
      : `https://huggingface.co/${DSPARK.repo}/resolve/${DSPARK.revision}/${DSPARK.file}`;
  }
  assets(settings) {
    const files = this.urls.map((url, index) => ({
      url,
      size: MODEL.shardSizes[index],
      type: "model",
    }));
    if (settings.dspark !== "off")
      files.push({ url: this.draftURL, size: DSPARK.bytes, type: "draft" });
    return files;
  }
  async cacheState(settings = this.settings || DEFAULTS) {
    const entries = await this.cache.list(),
      files = this.assets(settings),
      complete = files.map((f) =>
        entries.find(
          (e) =>
            e.metadata?.originalURL === f.url &&
            e.size === f.size &&
            e.metadata.originalSize === f.size,
        ),
      );
    return {
      entries,
      complete,
      cached: complete.every(Boolean),
      bytes: complete.reduce((n, f) => n + (f?.size || 0), 0),
    };
  }
  async load(settings) {
    if (this.loading) throw new Error("Model is already loading.");
    this.loading = true;
    this.loadController = new AbortController();
    const signal = this.loadController.signal,
      start = performance.now();
    this.logs = [];
    this.gpuLayers = 0;
    this.downloadMetrics = { requests: 0, retries: 0 };
    try {
      if (this.ready) await this.unload();
      const gpu = await this.probe();
      if (!gpu.ok) throw new Error(gpu.message);
      const files = this.assets(settings),
        total = files.reduce((n, f) => n + f.size, 0),
        cached = settings.saveModel
          ? await this.cacheState(settings)
          : { complete: [], cached: false };
      const blobs = await Promise.all(
        files.map(
          (file, i) =>
            this.sessionFiles.get(file.url) ||
            (cached.complete[i]
              ? this.cache.open(cached.complete[i].name)
              : null),
        ),
      );
      const restored = blobs.reduce((sum, blob, index) => {
        if (blob?.size !== files[index].size) blobs[index] = null;
        return sum + (blobs[index]?.size || 0);
      }, 0);
      this.onStatus({
        phase: restored === total ? "compile" : "download",
        progress: restored / total,
        cached: cached.cached,
        label:
          restored === total
            ? "Restoring this model"
            : "Downloading Q8 model" +
              (settings.dspark !== "off" ? " + DSpark" : ""),
      });
      if (settings.saveModel) navigator.storage?.persist?.().catch(() => {});
      const missing = files
        .map((file, index) => ({ ...file, index }))
        .filter((f) => !blobs[f.index]);
      if (missing.length)
        await downloadForSession(
          missing.map((f) => f.url),
          {
            signal,
            total,
            fileSizes: missing.map((f) => f.size),
            parallel: gpu.mobile ? 4 : 8,
            onProgress: ({ loaded, requests, retries }) => {
              this.downloadMetrics = { requests, retries };
              this.onStatus({
                phase: "download",
                progress: (loaded + restored) / total,
                loaded: loaded + restored,
                total,
                elapsed: (performance.now() - start) / 1000,
                label:
                  settings.dspark !== "off"
                    ? "Downloading Q8 model + DSpark"
                    : "Downloading Q8 model",
                persistent: settings.saveModel,
              });
            },
            onFileComplete: async (blob, index) => {
              const file = missing[index];
              if (
                file.type === "model" &&
                new TextDecoder().decode(
                  await blob.slice(0, 4).arrayBuffer(),
                ) !== "GGUF"
              )
                throw new Error(
                  "The weight download was not a valid GGUF file. Please retry.",
                );
              this.sessionFiles.set(file.url, blob);
              blobs[file.index] = blob;
              if (settings.saveModel) {
                signal.throwIfAborted();
                await this.cache.write(
                  await this.cache.getNameFromURL(file.url),
                  blob.stream(),
                  { originalURL: file.url, originalSize: blob.size, etag: "" },
                );
              }
            },
          },
        );
      signal.throwIfAborted();
      this.downloadSeconds = (performance.now() - start) / 1000;
      this.onStatus({
        phase: "compile",
        progress: 1,
        label: "Preparing dedicated WebGPU kernels",
      });
      const report = (event) => {
        if (event.diagnostic) {
          this.logs.push({
            level: event.level || "info",
            line: event.diagnostic,
          });
          if (this.logs.length > 100) this.logs.shift();
        }
        this.onStatus(event);
      };
      const { MiniCPMRuntime } = await import("./minicpm-runtime.js");
      this.dedicated = new MiniCPMRuntime(report);
      const runtime = this.dedicated;
      await startRuntime(
        async () => {
          await runtime.load(blobs.slice(0, 6), settings, signal);
          if (settings.dspark !== "off") {
            report({
              diagnostic:
                "Preparing DSpark draft model and target verification",
              level: "info",
            });
            const { DSparkRuntime } = await import("./dspark-runtime.js");
            const draft = new DSparkRuntime(report);
            runtime.draft = draft;
            await draft.initialize(runtime, blobs[6], signal);
          }
        },
        { signal, stop: () => runtime.exit(), timeoutMs: 120000 },
      );
      this.info = {
        n_ctx: settings.context,
        n_layer: 42,
        n_embd: 2048,
        n_vocab: 130560,
      };
      this.gpuLayers = 43;
      this.ready = true;
      this.settings = { ...settings };
      this.sessionFiles.clear();
      this.loadSeconds = (performance.now() - start) / 1000;
      this.gpuSeconds = this.loadSeconds - this.downloadSeconds;
      this.runtimeName = "MiniCPM5 dedicated WebGPU";
      this.runtimeCapabilities = {
        subgroups: runtime.subgroups,
        speculative: !!runtime.draft,
      };
      report({
        diagnostic: `Dedicated WebGPU ready: ${runtime.subgroups ? "subgroup" : "portable"} kernels${runtime.draft ? ", DSpark enabled" : ""}. All 43 layers on the GPU.`,
        level: "info",
      });
      this.onStatus({
        phase: "ready",
        label: "Model ready · dedicated WebGPU",
        seconds: this.loadSeconds,
        cached: cached.cached,
        info: this.info,
      });
    } catch (error) {
      signal.aborted || this.loadController.abort();
      await this.dedicated?.exit().catch(() => {});
      this.dedicated = null;
      this.ready = false;
      throw error;
    } finally {
      this.loading = false;
    }
  }
  async chat(messages, settings, tools, signal, onChunk) {
    if (!this.ready) throw new Error("Load the model first.");
    const prompt = this.dedicated.renderPrompt(messages, settings, tools),
      splitter = { thinking: settings.thinking, pending: "" };
    let last;
    await this.dedicated.generate(prompt, settings, signal, async (token) => {
      if (token.progress) {
        await onChunk({ choices: [], prompt_progress: token.progress });
        return;
      }
      last = token;
      await onChunk({
        choices: [
          {
            delta: splitStreamContent(splitter, token.delta),
            finish_reason: token.finish,
          },
        ],
        timings: {
          predicted_n: token.index,
          predicted_per_second:
            token.index > 1
              ? (token.index - 1) /
                Math.max(0.001, (token.elapsed - token.prefillMs) / 1000)
              : 0,
          prompt_n: token.promptTokens,
          prompt_ms: token.prefillMs,
        },
      });
    });
    if (last)
      await onChunk({
        choices: [
          {
            delta: splitStreamContent(splitter, "", true),
            finish_reason: last.finish || "stop",
          },
        ],
        usage: { completion_tokens: last.index },
        speculation: this.dedicated.specStats,
      });
  }
  async unload() {
    try {
      await this.dedicated?.exit();
    } finally {
      this.ready = false;
      this.dedicated = null;
    }
  }
  async clearCache() {
    await this.cache.deleteMany(
      (f) =>
        f.metadata?.originalURL?.includes(MODEL.weights) ||
        f.metadata?.originalURL?.includes(DSPARK.repo) ||
        f.name.includes("MiniCPM5-2B-Q8_0") ||
        f.metadata?.originalURL?.endsWith("/__qa/dspark"),
    );
  }
}