/**
 * CLI subprocess pre-warming.
 *
 * The Agent SDK's `startup()` spawns the Claude Code subprocess and completes
 * its initialize handshake ahead of time, so the first `query()` writes the
 * prompt directly to a ready process (~20× faster first response).
 *
 * Trade-off: all options (model, systemPrompt, mcpServers, agents, env, cwd)
 * are baked into the warm query at startup time. If the caller's options
 * don't match, the warm query can't be used and we fall back to a cold start.
 */

import { startup, type WarmQuery, type Options } from '@anthropic-ai/claude-agent-sdk';
import crypto from 'crypto';
import { log } from '../shared/logger.js';

interface CachedWarmup {
  key: string;
  warmQuery: WarmQuery;
  /** The abortController baked into the warm subprocess — caller must reuse this
   *  if it wants to abort the query. */
  abortController: AbortController;
}

let cached: CachedWarmup | null = null;
let inflight: Promise<void> | null = null;

/**
 * Build a cache key from the options that must match between preWarm() and
 * claimWarmup(). Lived options (like per-turn abortController or stderr
 * callbacks) are excluded — the SDK wires those at startup time but they're
 * not relevant to compatibility.
 */
function keyFor(options: Options): string {
  const keyable = {
    cwd: options.cwd,
    model: options.model,
    permissionMode: options.permissionMode,
    systemPrompt: options.systemPrompt,
    mcpServers: options.mcpServers
      ? Object.keys(options.mcpServers).sort().map((k) => [k, options.mcpServers![k]])
      : null,
    agents: options.agents
      ? Object.keys(options.agents).sort().map((k) => [k, options.agents![k]])
      : null,
    env: options.env
      ? Object.keys(options.env).sort().filter((k) => !k.startsWith('npm_')).map((k) => [k, options.env![k]])
      : null,
    resume: options.resume,
    betas: options.betas,
  };
  return crypto.createHash('sha256').update(JSON.stringify(keyable)).digest('hex');
}

export interface ClaimedWarmup {
  warmQuery: WarmQuery;
  abortController: AbortController;
}

/**
 * Spawn a pre-warmed subprocess with the given options. Fire-and-forget: if
 * preWarm is already in flight or the cached warmup already matches, no-op.
 *
 * We own the abortController so the caller can reuse it after claiming —
 * otherwise `.abort()` wouldn't reach the warm subprocess.
 */
export async function preWarm(options: Omit<Options, 'abortController'>): Promise<void> {
  if (inflight) return inflight;
  const key = keyFor(options as Options);
  if (cached?.key === key) return;

  inflight = (async () => {
    try {
      if (cached && cached.key !== key) {
        try { cached.warmQuery.close(); } catch {}
        cached = null;
      }
      const abortController = new AbortController();
      log.info('[cli-warmup] Pre-warming Claude subprocess...');
      const warmQuery = await startup({ options: { ...options, abortController } });
      cached = { key, warmQuery, abortController };
      log.ok('[cli-warmup] Subprocess pre-warmed');
    } catch (err: any) {
      log.warn(`[cli-warmup] Pre-warm failed: ${err?.message || err}`);
    } finally {
      inflight = null;
    }
  })();
  return inflight;
}

/**
 * Atomically claim the warm query if its options match. Returns null if
 * there's no warmup or the options differ — caller should cold-start.
 *
 * The caller must use the returned `abortController` to abort — the one baked
 * into the subprocess is the only one that works.
 */
export function claimWarmup(options: Omit<Options, 'abortController'>): ClaimedWarmup | null {
  if (!cached) return null;
  if (cached.key !== keyFor(options as Options)) return null;
  const claimed: ClaimedWarmup = { warmQuery: cached.warmQuery, abortController: cached.abortController };
  cached = null;
  log.info('[cli-warmup] Claimed pre-warmed subprocess');
  return claimed;
}

/** Close and discard any pending warmup (e.g. on shutdown or auth change). */
export function discardWarmup(): void {
  if (cached) {
    try { cached.warmQuery.close(); } catch {}
    cached = null;
  }
}