/**
 * Live benchmark runner. Requires LM Studio running locally with a model
 * loaded. Reuses the plugin's own modules end to end: token counting,
 * planning, chunk summarization (incl. transient retry), consolidation,
 * and content-addressed caching.
 *
 *   npm run bench
 */
import { Chat, LLM, LMStudioClient } from "@lmstudio/sdk";
import { mkdir, writeFile } from "fs/promises";
import { join } from "path";
import { ChunkCache } from "../src/cache";
import { pickConsolidation } from "../src/consolidate";
import { SummarizeCtl, summarizeChunk } from "../src/handler";
import { planCompaction } from "../src/plan";
import { noThinkDirective } from "../src/reasoning";
import {
  PromptVariant,
  buildMergePrompt,
  composeSystemMessage,
  stripReasoning,
  summarySystemPrompt,
} from "../src/summarizer";
import { TokenCounter } from "../src/tokens";
import {
  SYNTHETIC_USER_MESSAGE,
  CanonMessage,
  hashNamespace,
  hasUserMessage,
  prefixHashes,
} from "../src/view";
import { Fixture, allFixtures } from "./fixtures";
import {
  checkViewStructure,
  gradeAnswer,
  latencyStats,
  reductionPercent,
} from "./metrics";
import { BenchRun, FixtureResult, renderMarkdown, resultFileBase } from "./report";

const SETTINGS = {
  limitTokens: 500,
  floorTokens: 400,
  keepRecentTokens: 300,
  chunkTokens: 800,
  summarizerMaxTokens: 400,
  consolidationBudgetTokens: 600,
  consolidationKeepNewest: 2,
  mergeMaxTokens: 400,
  maxExcerptChars: 32000,
  quizMaxTokens: 120,
};

function makeBenchCtl(signal: AbortSignal): SummarizeCtl {
  return {
    abortSignal: signal,
    onAborted(listener: () => void) {
      signal.addEventListener("abort", listener, { once: true });
    },
  };
}

/**
 * Minimal `--variant full|compact` flag (default "full"), so
 * `npm run bench -- --variant compact` benches the compact summarizer prompt
 * against the same fixtures. Anything else on argv is ignored.
 */
function parseArgs(argv: string[]): { variant: PromptVariant } {
  let variant: PromptVariant = "full";
  for (let i = 0; i < argv.length; i++) {
    if (argv[i] !== "--variant") continue;
    const value = argv[i + 1];
    if (value !== "full" && value !== "compact") {
      console.error(
        `Invalid --variant value: ${value ?? "(missing)"} (expected "full" or "compact")`,
      );
      process.exit(1);
    }
    variant = value;
    i++;
  }
  return { variant };
}

async function quiz(
  model: LLM,
  systemText: string,
  question: string,
): Promise<string> {
  const result = await model.respond(
    Chat.from([
      { role: "system", content: systemText },
      { role: "user", content: question },
    ]),
    {
      maxTokens: SETTINGS.quizMaxTokens,
      temperature: 0,
      contextOverflowPolicy: "stopAtLimit",
    },
  );
  return stripReasoning(result.content);
}

async function runFixture(
  model: LLM,
  fixture: Fixture,
  variant: PromptVariant,
): Promise<FixtureResult> {
  const ctl = makeBenchCtl(new AbortController().signal);
  const counter = new TokenCounter((text) => model.countTokens(text));
  const tokens = await counter.countMessages(fixture.msgs);
  // Mirrors src/handler.ts's hashNamespace call exactly (schema 3, the
  // (variant, agentic) prompt identity, and the resolved system/foldSystem
  // prompt shape) so bench never silently drifts from the real cache
  // namespace. Bench always renders summarizer prompts in "system" shape
  // (separate system+user messages, never folded) — see
  // src/viewPlan.ts#promptChatMessages — so promptShape is hardcoded here.
  const namespace = hashNamespace({
    schema: 3,
    summarizerPromptCore: summarySystemPrompt({ variant, agentic: false }),
    agenticAddendumPrompt: summarySystemPrompt({ variant, agentic: true }),
    promptVariant: variant,
    promptShape: "system",
    model: model.identifier ?? "bench",
    chunkTokens: SETTINGS.chunkTokens,
    summarizerMaxTokens: SETTINGS.summarizerMaxTokens,
  });
  const hashes = prefixHashes(fixture.msgs, namespace);

  const plan = planCompaction({
    msgs: fixture.msgs,
    tokens,
    coveredUpTo: 0,
    summaryTokens: 0,
    estSummaryTokensPerChunk: SETTINGS.summarizerMaxTokens,
    limit: SETTINGS.limitTokens,
    floor: SETTINGS.floorTokens,
    autoCompact: true,
    force: true,
    keepRecentTokens: SETTINGS.keepRecentTokens,
    chunkTokens: SETTINGS.chunkTokens,
  });
  if (plan.cuts.length === 0) {
    throw new Error(`${fixture.name}: no compaction cuts produced`);
  }

  const cache = new ChunkCache();
  const chunkSummaries: string[] = [];
  const latencies: number[] = [];
  let from = 0;
  for (const cut of plan.cuts) {
    const started = Date.now();
    const summary = await summarizeChunk(
      ctl,
      model,
      fixture.msgs.slice(from, cut),
      SETTINGS.summarizerMaxTokens,
      SETTINGS.maxExcerptChars,
      undefined, // onProgress
      undefined, // attachmentMemories
      "system", // promptStrategy — bench always uses system-shape prompts
      variant,
    );
    latencies.push(Date.now() - started);
    chunkSummaries.push(summary);
    cache.put(hashes[cut - 1], {
      chunkSummaries: [...chunkSummaries],
      coveredCount: cut,
      tokensBefore: tokens.slice(0, cut).reduce((a, b) => a + b, 0),
      createdAt: Date.now(),
      lastUsedAt: Date.now(),
    });
    from = cut;
  }
  const coveredUpTo = from;
  const tail = fixture.msgs.slice(coveredUpTo);
  const systemText =
    composeSystemMessage([], chunkSummaries) ?? "(no summaries)";
  const tailForCheck: CanonMessage[] = hasUserMessage(tail)
    ? tail
    : [{ role: "user", text: SYNTHETIC_USER_MESSAGE }, ...tail];
  const structural = checkViewStructure(systemText, tailForCheck);

  let correctAfterCompaction = 0;
  for (const fact of fixture.facts) {
    const answer = await quiz(model, systemText, fact.question);
    if (gradeAnswer(answer, fact.accepted)) correctAfterCompaction++;
  }

  // Forced consolidation → drift measurement
  let correctAfterConsolidation: number | undefined;
  const summaryCounts = await Promise.all(
    chunkSummaries.map((s) => model.countTokens(s)),
  );
  const pick = pickConsolidation(
    chunkSummaries,
    summaryCounts,
    SETTINGS.consolidationBudgetTokens,
    SETTINGS.consolidationKeepNewest,
    SETTINGS.mergeMaxTokens,
  );
  if (pick !== null) {
    const directive = noThinkDirective(model.identifier);
    const prompt = buildMergePrompt(pick.merge, SETTINGS.maxExcerptChars, {
      noThinkDirective: directive,
      variant,
    });
    const merged = stripReasoning(
      (
        await model.respond(
          Chat.from([
            { role: "system", content: prompt.system },
            { role: "user", content: prompt.user },
          ]),
          {
            maxTokens: SETTINGS.mergeMaxTokens,
            temperature: 0.2,
            contextOverflowPolicy: "stopAtLimit",
          },
        )
      ).content,
    );
    const consolidated = [merged, ...pick.keep];
    const consolidatedSystem =
      composeSystemMessage([], consolidated) ?? "(no summaries)";
    correctAfterConsolidation = 0;
    for (const fact of fixture.facts) {
      const answer = await quiz(model, consolidatedSystem, fact.question);
      if (gradeAnswer(answer, fact.accepted)) correctAfterConsolidation++;
    }
  }

  const beforeTokens = tokens.reduce((a, b) => a + b, 0);
  const afterTokens =
    (await model.countTokens(systemText)) +
    tokens.slice(coveredUpTo).reduce((a, b) => a + b, 0);

  // Cache determinism: recomputing hashes must reuse every cached chunk with
  // zero further model calls.
  const rerunHashes = prefixHashes(fixture.msgs, namespace);
  const match = cache.findLongestMatch(rerunHashes);
  const cacheDeterministic =
    match !== undefined &&
    match.index === coveredUpTo - 1 &&
    match.entry.chunkSummaries.length === chunkSummaries.length;

  return {
    fixture: fixture.name,
    asked: fixture.facts.length,
    correctAfterCompaction,
    correctAfterConsolidation,
    reductionPercent: reductionPercent(beforeTokens, afterTokens),
    latency: latencyStats(latencies),
    structural,
    cacheDeterministic,
  };
}

async function main(): Promise<void> {
  const client = new LMStudioClient();
  let model: LLM;
  try {
    model = (await Promise.race([
      client.llm.model(),
      new Promise<never>((_, reject) =>
        setTimeout(() => reject(new Error("timeout")), 8000),
      ),
    ])) as LLM;
  } catch {
    console.error(
      "Could not reach LM Studio with a loaded model.\n" +
        "Start LM Studio, run `lms server start`, load a model, then retry.",
    );
    process.exit(1);
  }
  const info = await model.getModelInfo().catch(() => undefined);
  const modelId = info?.identifier ?? model.identifier ?? "unknown-model";
  const { variant } = parseArgs(process.argv.slice(2));
  console.log(
    `Benchmarking against ${modelId}${variant !== "full" ? ` (variant: ${variant})` : ""}…`,
  );

  const results: FixtureResult[] = [];
  for (const fixture of allFixtures()) {
    console.log(`  ${fixture.name}…`);
    results.push(await runFixture(model, fixture, variant));
  }

  const run: BenchRun = {
    date: new Date().toISOString().slice(0, 10),
    modelId,
    variant,
    settings: SETTINGS,
    fixtures: results,
  };
  const outDir = join(__dirname, "results");
  await mkdir(outDir, { recursive: true });
  const base = resultFileBase(run);
  await writeFile(join(outDir, `${base}.json`), JSON.stringify(run, null, 2));
  await writeFile(join(outDir, `${base}.md`), renderMarkdown(run));
  console.log(renderMarkdown(run));
  console.log(`\nWritten to bench/results/${base}.{json,md}`);
}

void main();

/**
 * Live benchmark runner. Requires LM Studio running locally with a model
 * loaded. Reuses the plugin's own modules end to end: token counting,
 * planning, chunk summarization (incl. transient retry), consolidation,
 * and content-addressed caching.
 *
 *   npm run bench
 */
import { Chat, LLM, LMStudioClient } from "@lmstudio/sdk";
import { mkdir, writeFile } from "fs/promises";
import { join } from "path";
import { ChunkCache } from "../src/cache";
import { pickConsolidation } from "../src/consolidate";
import { SummarizeCtl, summarizeChunk } from "../src/handler";
import { planCompaction } from "../src/plan";
import { noThinkDirective } from "../src/reasoning";
import {
  PromptVariant,
  buildMergePrompt,
  composeSystemMessage,
  stripReasoning,
  summarySystemPrompt,
} from "../src/summarizer";
import { TokenCounter } from "../src/tokens";
import {
  SYNTHETIC_USER_MESSAGE,
  CanonMessage,
  hashNamespace,
  hasUserMessage,
  prefixHashes,
} from "../src/view";
import { Fixture, allFixtures } from "./fixtures";
import {
  checkViewStructure,
  gradeAnswer,
  latencyStats,
  reductionPercent,
} from "./metrics";
import { BenchRun, FixtureResult, renderMarkdown, resultFileBase } from "./report";

const SETTINGS = {
  limitTokens: 500,
  floorTokens: 400,
  keepRecentTokens: 300,
  chunkTokens: 800,
  summarizerMaxTokens: 400,
  consolidationBudgetTokens: 600,
  consolidationKeepNewest: 2,
  mergeMaxTokens: 400,
  maxExcerptChars: 32000,
  quizMaxTokens: 120,
};

function makeBenchCtl(signal: AbortSignal): SummarizeCtl {
  return {
    abortSignal: signal,
    onAborted(listener: () => void) {
      signal.addEventListener("abort", listener, { once: true });
    },
  };
}

/**
 * Minimal `--variant full|compact` flag (default "full"), so
 * `npm run bench -- --variant compact` benches the compact summarizer prompt
 * against the same fixtures. Anything else on argv is ignored.
 */
function parseArgs(argv: string[]): { variant: PromptVariant } {
  let variant: PromptVariant = "full";
  for (let i = 0; i < argv.length; i++) {
    if (argv[i] !== "--variant") continue;
    const value = argv[i + 1];
    if (value !== "full" && value !== "compact") {
      console.error(
        `Invalid --variant value: ${value ?? "(missing)"} (expected "full" or "compact")`,
      );
      process.exit(1);
    }
    variant = value;
    i++;
  }
  return { variant };
}

async function quiz(
  model: LLM,
  systemText: string,
  question: string,
): Promise<string> {
  const result = await model.respond(
    Chat.from([
      { role: "system", content: systemText },
      { role: "user", content: question },
    ]),
    {
      maxTokens: SETTINGS.quizMaxTokens,
      temperature: 0,
      contextOverflowPolicy: "stopAtLimit",
    },
  );
  return stripReasoning(result.content);
}

async function runFixture(
  model: LLM,
  fixture: Fixture,
  variant: PromptVariant,
): Promise<FixtureResult> {
  const ctl = makeBenchCtl(new AbortController().signal);
  const counter = new TokenCounter((text) => model.countTokens(text));
  const tokens = await counter.countMessages(fixture.msgs);
  // Mirrors src/handler.ts's hashNamespace call exactly (schema 3, the
  // (variant, agentic) prompt identity, and the resolved system/foldSystem
  // prompt shape) so bench never silently drifts from the real cache
  // namespace. Bench always renders summarizer prompts in "system" shape
  // (separate system+user messages, never folded) — see
  // src/viewPlan.ts#promptChatMessages — so promptShape is hardcoded here.
  const namespace = hashNamespace({
    schema: 3,
    summarizerPromptCore: summarySystemPrompt({ variant, agentic: false }),
    agenticAddendumPrompt: summarySystemPrompt({ variant, agentic: true }),
    promptVariant: variant,
    promptShape: "system",
    model: model.identifier ?? "bench",
    chunkTokens: SETTINGS.chunkTokens,
    summarizerMaxTokens: SETTINGS.summarizerMaxTokens,
  });
  const hashes = prefixHashes(fixture.msgs, namespace);

  const plan = planCompaction({
    msgs: fixture.msgs,
    tokens,
    coveredUpTo: 0,
    summaryTokens: 0,
    estSummaryTokensPerChunk: SETTINGS.summarizerMaxTokens,
    limit: SETTINGS.limitTokens,
    floor: SETTINGS.floorTokens,
    autoCompact: true,
    force: true,
    keepRecentTokens: SETTINGS.keepRecentTokens,
    chunkTokens: SETTINGS.chunkTokens,
  });
  if (plan.cuts.length === 0) {
    throw new Error(`${fixture.name}: no compaction cuts produced`);
  }

  const cache = new ChunkCache();
  const chunkSummaries: string[] = [];
  const latencies: number[] = [];
  let from = 0;
  for (const cut of plan.cuts) {
    const started = Date.now();
    const summary = await summarizeChunk(
      ctl,
      model,
      fixture.msgs.slice(from, cut),
      SETTINGS.summarizerMaxTokens,
      SETTINGS.maxExcerptChars,
      undefined, // onProgress
      undefined, // attachmentMemories
      "system", // promptStrategy — bench always uses system-shape prompts
      variant,
    );
    latencies.push(Date.now() - started);
    chunkSummaries.push(summary);
    cache.put(hashes[cut - 1], {
      chunkSummaries: [...chunkSummaries],
      coveredCount: cut,
      tokensBefore: tokens.slice(0, cut).reduce((a, b) => a + b, 0),
      createdAt: Date.now(),
      lastUsedAt: Date.now(),
    });
    from = cut;
  }
  const coveredUpTo = from;
  const tail = fixture.msgs.slice(coveredUpTo);
  const systemText =
    composeSystemMessage([], chunkSummaries) ?? "(no summaries)";
  const tailForCheck: CanonMessage[] = hasUserMessage(tail)
    ? tail
    : [{ role: "user", text: SYNTHETIC_USER_MESSAGE }, ...tail];
  const structural = checkViewStructure(systemText, tailForCheck);

  let correctAfterCompaction = 0;
  for (const fact of fixture.facts) {
    const answer = await quiz(model, systemText, fact.question);
    if (gradeAnswer(answer, fact.accepted)) correctAfterCompaction++;
  }

  // Forced consolidation → drift measurement
  let correctAfterConsolidation: number | undefined;
  const summaryCounts = await Promise.all(
    chunkSummaries.map((s) => model.countTokens(s)),
  );
  const pick = pickConsolidation(
    chunkSummaries,
    summaryCounts,
    SETTINGS.consolidationBudgetTokens,
    SETTINGS.consolidationKeepNewest,
    SETTINGS.mergeMaxTokens,
  );
  if (pick !== null) {
    const directive = noThinkDirective(model.identifier);
    const prompt = buildMergePrompt(pick.merge, SETTINGS.maxExcerptChars, {
      noThinkDirective: directive,
      variant,
    });
    const merged = stripReasoning(
      (
        await model.respond(
          Chat.from([
            { role: "system", content: prompt.system },
            { role: "user", content: prompt.user },
          ]),
          {
            maxTokens: SETTINGS.mergeMaxTokens,
            temperature: 0.2,
            contextOverflowPolicy: "stopAtLimit",
          },
        )
      ).content,
    );
    const consolidated = [merged, ...pick.keep];
    const consolidatedSystem =
      composeSystemMessage([], consolidated) ?? "(no summaries)";
    correctAfterConsolidation = 0;
    for (const fact of fixture.facts) {
      const answer = await quiz(model, consolidatedSystem, fact.question);
      if (gradeAnswer(answer, fact.accepted)) correctAfterConsolidation++;
    }
  }

  const beforeTokens = tokens.reduce((a, b) => a + b, 0);
  const afterTokens =
    (await model.countTokens(systemText)) +
    tokens.slice(coveredUpTo).reduce((a, b) => a + b, 0);

  // Cache determinism: recomputing hashes must reuse every cached chunk with
  // zero further model calls.
  const rerunHashes = prefixHashes(fixture.msgs, namespace);
  const match = cache.findLongestMatch(rerunHashes);
  const cacheDeterministic =
    match !== undefined &&
    match.index === coveredUpTo - 1 &&
    match.entry.chunkSummaries.length === chunkSummaries.length;

  return {
    fixture: fixture.name,
    asked: fixture.facts.length,
    correctAfterCompaction,
    correctAfterConsolidation,
    reductionPercent: reductionPercent(beforeTokens, afterTokens),
    latency: latencyStats(latencies),
    structural,
    cacheDeterministic,
  };
}

async function main(): Promise<void> {
  const client = new LMStudioClient();
  let model: LLM;
  try {
    model = (await Promise.race([
      client.llm.model(),
      new Promise<never>((_, reject) =>
        setTimeout(() => reject(new Error("timeout")), 8000),
      ),
    ])) as LLM;
  } catch {
    console.error(
      "Could not reach LM Studio with a loaded model.\n" +
        "Start LM Studio, run `lms server start`, load a model, then retry.",
    );
    process.exit(1);
  }
  const info = await model.getModelInfo().catch(() => undefined);
  const modelId = info?.identifier ?? model.identifier ?? "unknown-model";
  const { variant } = parseArgs(process.argv.slice(2));
  console.log(
    `Benchmarking against ${modelId}${variant !== "full" ? ` (variant: ${variant})` : ""}…`,
  );

  const results: FixtureResult[] = [];
  for (const fixture of allFixtures()) {
    console.log(`  ${fixture.name}…`);
    results.push(await runFixture(model, fixture, variant));
  }

  const run: BenchRun = {
    date: new Date().toISOString().slice(0, 10),
    modelId,
    variant,
    settings: SETTINGS,
    fixtures: results,
  };
  const outDir = join(__dirname, "results");
  await mkdir(outDir, { recursive: true });
  const base = resultFileBase(run);
  await writeFile(join(outDir, `${base}.json`), JSON.stringify(run, null, 2));
  await writeFile(join(outDir, `${base}.md`), renderMarkdown(run));
  console.log(renderMarkdown(run));
  console.log(`\nWritten to bench/results/${base}.{json,md}`);
}

void main();