/**
 * Downloads and manages the local embedding runtime post-hatch.
 *
 * Instead of shipping heavy native + JS dependencies inside the .app bundle,
 * we download them from npm and run embeddings in a **separate bun process**.
 * The compiled daemon binary cannot resolve bare specifiers in dynamically
 * imported files, so we spawn a standalone bun process that runs an embed
 * worker script communicating via JSON-lines over stdin/stdout.
 *
 * The runtime is stored in $VELLUM_WORKSPACE_DIR/embedding-models/ and used
 * by embedding-local.ts on demand.
 *
 * Follows the same download/install pattern as qdrant-manager.ts.
 */

import {
  existsSync,
  mkdirSync,
  readdirSync,
  readFileSync,
  renameSync,
  rmSync,
  writeFileSync,
} from "node:fs";
import { arch, platform } from "node:os";
import { join } from "node:path";

import { ensureBun, findBun } from "../../util/bun-runtime.js";
import { getLogger } from "../../util/logger.js";
import { getEmbeddingModelsDir } from "../../util/platform.js";
import { PromiseGuard } from "../../util/promise-guard.js";

const log = getLogger("embedding-runtime-manager");

// Pinned versions matching assistant/bun.lock
const ONNXRUNTIME_NODE_VERSION = "1.21.0";
const ONNXRUNTIME_COMMON_VERSION = "1.21.0";
const TRANSFORMERS_VERSION = "3.8.1";
const JINJA_VERSION = "0.5.5";

/**
 * Composite version string for cache invalidation. Bumping the trailing
 * `_workers-vN` suffix forces existing installs to regenerate the worker
 * scripts when the worker IPC contract or spawn-args list changes (without
 * requiring an `@huggingface/transformers` version bump).
 *
 * v4: workers read `VELLUM_ONNX_INTRA_OP_THREADS` and pass
 * `session_options.intraOpNumThreads` (JARVIS-1398). The environment the host
 * spawns them with counts as part of the contract, so an install that keeps
 * the v3 scripts would silently ignore the cap.
 */
export const RUNTIME_VERSION = `ort-${ONNXRUNTIME_NODE_VERSION}_hf-${TRANSFORMERS_VERSION}_jinja-${JINJA_VERSION}_workers-v4`;

const WORKER_FILENAME = "embed-worker.mjs";
const RERANK_WORKER_FILENAME = "rerank-worker.mjs";

/** Module-level guard so concurrent in-process calls share one download. */
const installGuard = new PromiseGuard<void>();

interface VersionManifest {
  runtimeVersion: string;
  onnxruntimeNodeVersion: string;
  onnxruntimeCommonVersion: string;
  transformersVersion: string;
  platform: string;
  arch: string;
  installedAt: string;
}

// ── npm tarball helpers ─────────────────────────────────────────────

function npmTarballUrl(pkg: string, version: string): string {
  // Scoped packages encode the scope in the URL
  const encoded = pkg.replace("/", "%2f");
  const basename = pkg.startsWith("@") ? pkg.split("/")[1] : pkg;
  return `https://registry.npmjs.org/${encoded}/-/${basename}-${version}.tgz`;
}

async function downloadAndExtract(
  url: string,
  targetDir: string,
  signal?: AbortSignal,
): Promise<void> {
  log.info({ url, targetDir }, "Downloading npm package");

  const response = await fetch(url, { signal });
  if (!response.ok) {
    throw new Error(
      `Failed to download ${url}: ${response.status} ${response.statusText}`,
    );
  }

  const tarball = await response.arrayBuffer();

  // npm tarballs extract to package/, we need to redirect to targetDir
  mkdirSync(targetDir, { recursive: true });

  const tmpTar = join(targetDir, `download-${Date.now()}.tgz`);
  writeFileSync(tmpTar, Buffer.from(tarball));

  try {
    // Extract tarball, stripping the leading "package/" directory
    const proc = Bun.spawn({
      cmd: ["tar", "xzf", tmpTar, "-C", targetDir, "--strip-components=1"],
      stdout: "ignore",
      stderr: "pipe",
    });
    await proc.exited;
    if (proc.exitCode !== 0) {
      const stderr = await new Response(proc.stderr).text();
      throw new Error(`Failed to extract ${url}: ${stderr}`);
    }
  } finally {
    try {
      rmSync(tmpTar);
    } catch {
      /* ignore */
    }
  }
}

// ── Worker script content ───────────────────────────────────────────

export function generateWorkerScript(): string {
  // This script is run by a standalone bun process (not the compiled daemon).
  // Because it runs in a real bun runtime, bare specifier resolution works
  // normally — node_modules/ in the same directory is found automatically.
  return `\
// embed-worker.mjs — Auto-generated by EmbeddingRuntimeManager
// Runs in a separate bun process, communicates via JSON-lines over stdin/stdout.
process.title = 'embed-worker';
import { pipeline, env } from '@huggingface/transformers';

const model = process.argv[2];
const cacheDir = process.argv[3];
if (cacheDir && env) env.cacheDir = cacheDir;

// Cap the ONNX intra-op thread pool. Left unset, ONNX takes one thread per
// physical core for the whole batch and starves foreground work; the host
// computes the cap in util/worker-compute.ts and passes it in the environment.
// A missing or unparseable value means no cap, i.e. the ONNX default.
const intraOpNumThreads = Number(process.env.VELLUM_ONNX_INTRA_OP_THREADS);
const sessionOptions =
  Number.isInteger(intraOpNumThreads) && intraOpNumThreads > 0
    ? { intraOpNumThreads }
    : undefined;

let extractor;
try {
  extractor = await pipeline('feature-extraction', model, { dtype: 'fp32', session_options: sessionOptions });
  process.stdout.write(JSON.stringify({ type: 'ready' }) + '\\n');
} catch (err) {
  process.stdout.write(JSON.stringify({ type: 'error', error: err.message || String(err) }) + '\\n');
  process.exit(1);
}

// Sequential request queue to avoid concurrent ONNX inference
const decoder = new TextDecoder();
let buffer = '';
let processing = false;
const queue = [];

process.stdin.on('data', (chunk) => {
  buffer += typeof chunk === 'string' ? chunk : decoder.decode(chunk, { stream: true });
  let idx;
  while ((idx = buffer.indexOf('\\n')) !== -1) {
    const line = buffer.slice(0, idx);
    buffer = buffer.slice(idx + 1);
    if (line.trim()) queue.push(line);
  }
  if (!processing) processQueue();
});

async function processQueue() {
  processing = true;
  while (queue.length > 0) {
    const line = queue.shift();
    let req;
    try { req = JSON.parse(line); } catch { continue; }
    try {
      const output = await extractor(req.texts, { pooling: 'cls', normalize: true });
      const vectors = output.tolist();
      process.stdout.write(JSON.stringify({ id: req.id, vectors }) + '\\n');
    } catch (err) {
      process.stdout.write(JSON.stringify({ id: req.id, error: err.message || String(err) }) + '\\n');
    }
  }
  processing = false;
}

process.stdin.on('end', () => process.exit(0));
`;
}

export function generateRerankWorkerScript(): string {
  // Cross-encoder rerank worker. Loads a sequence-classification model and
  // scores paired (queries[i], passages[i]) tuples in one batched ONNX
  // inference call. Mirrors the embed worker's lifecycle (ready signal,
  // JSON-lines IPC, sequential queue) so LocalRerankBackend can reuse the
  // same supervisor pattern.
  //
  // Request shape: { id, queries: string[], passages: string[] } with
  // queries.length === passages.length. Each pair is one (query, passage)
  // tuple; multiple distinct queries can ride in a single batch so the
  // activation pipeline can score the user-channel and assistant-channel
  // queries against a shared candidate set in one tokenizer + ONNX call.
  return `\
// rerank-worker.mjs — Auto-generated by EmbeddingRuntimeManager
// Runs in a separate bun process, communicates via JSON-lines over stdin/stdout.
process.title = 'rerank-worker';
import {
  AutoModelForSequenceClassification,
  AutoTokenizer,
  env,
} from '@huggingface/transformers';

const model = process.argv[2];
const cacheDir = process.argv[3];
const dtype = process.argv[4] || 'q8';
if (cacheDir && env) env.cacheDir = cacheDir;

// Cap the ONNX intra-op thread pool. See the embed worker for the rationale;
// the host computes the cap in util/worker-compute.ts and passes it in the
// environment. A missing or unparseable value means no cap, i.e. the default.
const intraOpNumThreads = Number(process.env.VELLUM_ONNX_INTRA_OP_THREADS);
const sessionOptions =
  Number.isInteger(intraOpNumThreads) && intraOpNumThreads > 0
    ? { intraOpNumThreads }
    : undefined;

let tokenizer;
let session;
try {
  tokenizer = await AutoTokenizer.from_pretrained(model);
  session = await AutoModelForSequenceClassification.from_pretrained(model, { dtype, session_options: sessionOptions });
  process.stdout.write(JSON.stringify({ type: 'ready' }) + '\\n');
} catch (err) {
  process.stdout.write(JSON.stringify({ type: 'error', error: err.message || String(err) }) + '\\n');
  process.exit(1);
}

const sigmoid = (x) => 1 / (1 + Math.exp(-x));

const decoder = new TextDecoder();
let buffer = '';
let processing = false;
const queue = [];

process.stdin.on('data', (chunk) => {
  buffer += typeof chunk === 'string' ? chunk : decoder.decode(chunk, { stream: true });
  let idx;
  while ((idx = buffer.indexOf('\\n')) !== -1) {
    const line = buffer.slice(0, idx);
    buffer = buffer.slice(idx + 1);
    if (line.trim()) queue.push(line);
  }
  if (!processing) processQueue();
});

async function processQueue() {
  processing = true;
  while (queue.length > 0) {
    const line = queue.shift();
    let req;
    try { req = JSON.parse(line); } catch { continue; }
    try {
      const { id, queries, passages } = req;
      if (
        !Array.isArray(queries) || !Array.isArray(passages) ||
        queries.length !== passages.length || passages.length === 0
      ) {
        process.stdout.write(JSON.stringify({ id, scores: [] }) + '\\n');
        continue;
      }
      const inputs = await tokenizer(queries, {
        text_pair: passages,
        padding: true,
        truncation: true,
      });
      const out = await session(inputs);
      const logits = out.logits.data;
      const scores = new Array(passages.length);
      for (let i = 0; i < passages.length; i++) {
        scores[i] = sigmoid(Number(logits[i]));
      }
      process.stdout.write(JSON.stringify({ id, scores }) + '\\n');
    } catch (err) {
      process.stdout.write(JSON.stringify({ id: req?.id, error: err.message || String(err) }) + '\\n');
    }
  }
  processing = false;
}

process.stdin.on('end', () => process.exit(0));
`;
}

// ── Main manager ────────────────────────────────────────────────────

export class EmbeddingRuntimeManager {
  private readonly baseDir: string;

  constructor(baseDir?: string) {
    this.baseDir = baseDir ?? getEmbeddingModelsDir();
  }

  /** Check if the embedding runtime is installed and up-to-date. */
  isReady(): boolean {
    const manifest = this.readManifest();
    if (!manifest) {
      return false;
    }
    if (manifest.runtimeVersion !== RUNTIME_VERSION) {
      return false;
    }

    // Verify both worker scripts exist and a bun binary is available
    return (
      existsSync(this.getWorkerPath()) &&
      existsSync(this.getRerankWorkerPath()) &&
      this.getBunPath() !== undefined
    );
  }

  /** Path to the embed worker script. */
  getWorkerPath(): string {
    return join(this.baseDir, WORKER_FILENAME);
  }

  /** Path to the rerank worker script. */
  getRerankWorkerPath(): string {
    return join(this.baseDir, RERANK_WORKER_FILENAME);
  }

  /**
   * Find a usable bun binary.
   * Delegates to the shared bun-runtime helper, also checking
   * the legacy per-runtime bin/ directory for backwards compat.
   */
  getBunPath(): string | undefined {
    // Check per-runtime bin/ directory for backwards compat
    const legacyBun = join(this.baseDir, "bin", "bun");
    if (existsSync(legacyBun)) {
      return legacyBun;
    }

    return findBun();
  }

  /**
   * Download and install the embedding runtime if not already present.
   * Safe to call concurrently — in-process calls share one promise via
   * PromiseGuard, and cross-process calls are serialized via a lock file.
   */
  async ensureInstalled(signal?: AbortSignal): Promise<void> {
    if (this.isReady()) {
      return;
    }

    // Deduplicate concurrent in-process calls
    await installGuard.run(() => this.acquireLockAndInstall(signal));

    // If another process was downloading and we skipped, or if the download
    // somehow failed silently, reset the guard so we can retry next time.
    if (!this.isReady()) {
      installGuard.reset();
    }
  }

  private async acquireLockAndInstall(signal?: AbortSignal): Promise<void> {
    // Re-check after acquiring the in-process guard
    if (this.isReady()) {
      return;
    }

    // Cross-process lock to prevent duplicate downloads
    const lockPath = join(this.baseDir, ".downloading");
    if (existsSync(lockPath)) {
      try {
        const lockContent = readFileSync(lockPath, "utf-8").trim();
        const lockPid = parseInt(lockContent, 10);
        if (!isNaN(lockPid) && lockPid !== process.pid) {
          try {
            process.kill(lockPid, 0);
            log.info(
              { lockPid },
              "Another process is downloading the embedding runtime, skipping",
            );
            return;
          } catch {
            log.info({ lockPid }, "Cleaning up stale download lock");
          }
        }
      } catch {
        // Can't read lock file, proceed
      }
    }

    mkdirSync(this.baseDir, { recursive: true });

    // Write a .gitignore so the workspace git repo ignores this directory
    const gitignorePath = join(this.baseDir, ".gitignore");
    if (!existsSync(gitignorePath)) {
      writeFileSync(gitignorePath, "*\n!.gitignore\n");
    }

    writeFileSync(lockPath, String(process.pid));

    try {
      await this.install(signal);
    } finally {
      try {
        rmSync(lockPath);
      } catch {
        /* ignore */
      }
    }
  }

  private async install(signal?: AbortSignal): Promise<void> {
    const os = platform();
    const cpu = arch();
    log.info(
      { os, cpu, runtimeVersion: RUNTIME_VERSION },
      "Installing embedding runtime",
    );

    // Work in a temp directory for atomic install
    const tmpDir = join(this.baseDir, `.installing-${Date.now()}`);
    mkdirSync(tmpDir, { recursive: true });

    // Declared outside try so catch/finally can reference them for cleanup
    const modelCacheDir = join(this.baseDir, "model-cache");
    let tmpModelCache: string | null = null;

    try {
      // Step 1: Download npm packages in parallel
      const nodeModules = join(tmpDir, "node_modules");
      const downloads: Promise<void>[] = [
        downloadAndExtract(
          npmTarballUrl("onnxruntime-node", ONNXRUNTIME_NODE_VERSION),
          join(nodeModules, "onnxruntime-node"),
          signal,
        ),
        downloadAndExtract(
          npmTarballUrl("onnxruntime-common", ONNXRUNTIME_COMMON_VERSION),
          join(nodeModules, "onnxruntime-common"),
          signal,
        ),
        downloadAndExtract(
          npmTarballUrl("@huggingface/transformers", TRANSFORMERS_VERSION),
          join(nodeModules, "@huggingface", "transformers"),
          signal,
        ),
        downloadAndExtract(
          npmTarballUrl("@huggingface/jinja", JINJA_VERSION),
          join(nodeModules, "@huggingface", "jinja"),
          signal,
        ),
      ];

      await Promise.all(downloads);

      // Ensure bun is available (downloads to shared location if needed)
      await ensureBun();

      if (signal?.aborted) {
        throw new DOMException("Aborted", "AbortError");
      }

      log.info("npm packages downloaded, stripping non-platform binaries");

      // Step 2: Strip non-platform native binaries
      const onnxBinDir = join(
        nodeModules,
        "onnxruntime-node",
        "bin",
        "napi-v3",
      );
      if (existsSync(onnxBinDir)) {
        const entries = readdirSync(onnxBinDir);
        for (const entry of entries) {
          // Keep all darwin architectures (arm64 and x86_64) since uname -m
          // is unreliable under Rosetta (returns x86_64 on Apple Silicon)
          if (entry !== os) {
            rmSync(join(onnxBinDir, entry), { recursive: true, force: true });
          }
        }
      }

      // Strip non-runtime files to reduce disk usage.
      // Keep lib/ directories — they contain JS entry points needed for bare
      // specifier imports in the worker subprocess.
      const onnxNodeDir = join(nodeModules, "onnxruntime-node");
      rmSync(join(onnxNodeDir, "script"), { recursive: true, force: true });
      rmSync(join(onnxNodeDir, "README.md"), { force: true });
      rmSync(join(nodeModules, "onnxruntime-common", "README.md"), {
        force: true,
      });

      // Step 3: Create a stub "sharp" package so that the pre-built
      // transformers.node.mjs can import it without error. The bundle checks
      // `if (sharp)` at module initialization time — the stub must be truthy.
      // We only use text embeddings, never image processing.
      const sharpDir = join(nodeModules, "sharp");
      mkdirSync(sharpDir, { recursive: true });
      writeFileSync(
        join(sharpDir, "package.json"),
        '{"name":"sharp","version":"0.0.0","main":"index.js"}\n',
      );
      writeFileSync(
        join(sharpDir, "index.js"),
        [
          "// Stub: only text embeddings are used, no image processing.",
          "// Must be a truthy function so transformers.node.mjs initialization passes.",
          'function sharp() { throw new Error("sharp stub: image processing not available"); }',
          "sharp.format = {};",
          "module.exports = sharp;",
          "",
        ].join("\n"),
      );

      // Step 4: Write embed + rerank worker scripts
      writeFileSync(join(tmpDir, WORKER_FILENAME), generateWorkerScript());
      writeFileSync(
        join(tmpDir, RERANK_WORKER_FILENAME),
        generateRerankWorkerScript(),
      );

      // Step 5: Write version manifest
      const manifest: VersionManifest = {
        runtimeVersion: RUNTIME_VERSION,
        onnxruntimeNodeVersion: ONNXRUNTIME_NODE_VERSION,
        onnxruntimeCommonVersion: ONNXRUNTIME_COMMON_VERSION,
        transformersVersion: TRANSFORMERS_VERSION,
        platform: os,
        arch: cpu,
        installedAt: new Date().toISOString(),
      };
      writeFileSync(
        join(tmpDir, "version.json"),
        JSON.stringify(manifest, null, 2) + "\n",
      );

      // Step 6: Atomic swap — remove old install and rename temp to final
      // Preserve model-cache/ and .gitignore
      const hadModelCache = existsSync(modelCacheDir);
      if (hadModelCache) {
        tmpModelCache = join(
          this.baseDir,
          `.model-cache-preserve-${Date.now()}`,
        );
        renameSync(modelCacheDir, tmpModelCache);
      }

      // Remove old install (preserving dotfiles like .gitignore, .downloading, temp dirs)
      for (const entry of readdirSync(this.baseDir)) {
        if (entry.startsWith(".") || entry === tmpDir.split("/").pop()) {
          continue;
        }
        rmSync(join(this.baseDir, entry), { recursive: true, force: true });
      }

      // Move new files into place
      for (const entry of readdirSync(tmpDir)) {
        renameSync(join(tmpDir, entry), join(this.baseDir, entry));
      }

      // Restore model cache
      if (tmpModelCache && existsSync(tmpModelCache)) {
        renameSync(tmpModelCache, modelCacheDir);
      }

      log.info(
        { runtimeVersion: RUNTIME_VERSION },
        "Embedding runtime installed successfully",
      );
    } catch (err) {
      // Restore preserved directories if the swap failed
      if (
        tmpModelCache &&
        existsSync(tmpModelCache) &&
        !existsSync(modelCacheDir)
      ) {
        try {
          renameSync(tmpModelCache, modelCacheDir);
        } catch {
          /* best effort */
        }
      }
      log.error({ err }, "Failed to install embedding runtime");
      throw err;
    } finally {
      // Clean up temp directory and any leftover preserve dirs
      rmSync(tmpDir, { recursive: true, force: true });
      if (tmpModelCache) {
        rmSync(tmpModelCache, { recursive: true, force: true });
      }
    }
  }

  private readManifest(): VersionManifest | null {
    const manifestPath = join(this.baseDir, "version.json");
    if (!existsSync(manifestPath)) {
      return null;
    }
    try {
      return JSON.parse(readFileSync(manifestPath, "utf-8"));
    } catch {
      return null;
    }
  }
}