import { existsSync, mkdirSync, mkdtempSync, rmSync, symlinkSync } from "node:fs";
import { spawnSync } from "node:child_process";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { fileURLToPath, pathToFileURL } from "node:url";
import { resolvePrimeRoot } from "./prime-root.mjs";

const primeRoot = resolvePrimeRoot();
const prime = await import(pathToFileURL(join(primeRoot, "dist/index.js")));
const ai = await import(pathToFileURL(join(
  primeRoot,
  "node_modules/@earendil-works/pi-ai/dist/index.js",
)));
const root = mkdtempSync(join(tmpdir(), "prime-snapcompact-auto-"));
const sessionDir = join(root, "sessions");
const agentDir = join(root, "agent");
mkdirSync(join(agentDir, "extensions"), { recursive: true });
symlinkSync(
  fileURLToPath(new URL("..", import.meta.url)),
  join(agentDir, "extensions", "snapcompact"),
  "dir",
);
const registration = ai.registerFauxProvider({
  models: [{
    id: "snapcompact-vision",
    name: "SnapCompact vision test",
    input: ["text", "image"],
    reasoning: false,
    contextWindow: 100_000,
    maxTokens: 8192,
  }],
  tokensPerSecond: 100_000,
});
const model = registration.getModel();
const seenContexts = [];
const response = (text) => (context) => {
  seenContexts.push(context);
  return ai.fauxAssistantMessage(text);
};
registration.setResponses([
  response("first response"),
  response("second response"),
]);
const auth = prime.AuthStorage.inMemory();
auth.setRuntimeApiKey(model.provider, "snapcompact-test");
const modelRegistry = prime.ModelRegistry.inMemory(auth);
const cwd = process.cwd();
const sessionManager = prime.SessionManager.create(cwd, sessionDir);
const usage = {
  input: 94_000,
  output: 1000,
  cacheRead: 0,
  cacheWrite: 0,
  totalTokens: 95_000,
  cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
};
for (let turn = 0; turn < 8; turn++) {
  sessionManager.appendMessage({
    role: "user",
    content: `analyze repository turn ${turn}`,
    timestamp: Date.now(),
  });
  sessionManager.appendMessage({
    role: "assistant",
    content: [{
      type: "text",
      text: (`turn ${turn} repository analysis src/file.ts [] {}\n`).repeat(1800),
    }],
    api: model.api,
    provider: model.provider,
    model: model.id,
    usage,
    stopReason: "stop",
    timestamp: Date.now(),
  });
}
sessionManager.flushNow();

const sessionOptions = (manager) => ({
  cwd,
  agentDir,
  sessionManager: manager,
  model,
  modelRegistry,
  noTools: "all",
  includeGoals: false,
  includeCompactSkill: false,
  prewarmIpythonKernel: false,
  telemetryDisabled: true,
});
const imageCount = (context) => context.messages.reduce(
  (count, message) => count + (
    Array.isArray(message.content)
      ? message.content.filter((block) => block.type === "image").length
      : 0
  ),
  0,
);

let { session, extensionsResult } = await prime.createAgentSession(sessionOptions(sessionManager));
await session.bindExtensions({});
if (extensionsResult.errors.length) throw new Error(JSON.stringify(extensionsResult.errors));
try {
  await session.prompt("continue after automatic compaction");
  const entry = prime.getLatestCompactionEntry(sessionManager.getBranch());
  if (!entry?.fromHook || !entry.details?.snapcompact?.artifact) {
    throw new Error("automatic compaction did not persist SnapCompact details");
  }
  const firstImages = imageCount(seenContexts[0]);
  if (firstImages < 1 || firstImages > entry.details.snapcompact.frameBudget) {
    throw new Error("automatic post-compaction request did not receive bounded archive images");
  }
  const detailsBytes = Buffer.byteLength(JSON.stringify(entry.details));
  if (detailsBytes >= 1024 * 1024 || JSON.stringify(entry.details).includes("iVBORw0KGgo")) {
    throw new Error("automatic compaction inlined its visual archive");
  }
  const artifactPath = join(
    root,
    "session-artifacts",
    session.sessionId,
    "snapcompact",
    entry.details.snapcompact.artifact.file,
  );
  if (!existsSync(artifactPath)) throw new Error("automatic compaction artifact is missing");

  await session.prompt("one more request");
  const secondImages = imageCount(seenContexts[1]);
  if (secondImages < 1) throw new Error("repeated request lost the active visual archive");

  const sessionFile = session.sessionFile;
  await session.disposeAsync();
  registration.appendResponses([response("resumed response")]);
  const resumedManager = prime.SessionManager.open(sessionFile, sessionDir, cwd);
  ({ session, extensionsResult } = await prime.createAgentSession(sessionOptions(resumedManager)));
  await session.bindExtensions({});
  if (extensionsResult.errors.length) throw new Error(JSON.stringify(extensionsResult.errors));
  await session.prompt("continue the resumed archive");
  const resumedImages = imageCount(seenContexts[2]);
  if (resumedImages < 1) throw new Error("resumed session did not rehydrate archive images");

  const forkSourceFile = session.sessionFile;
  const forkFile = resumedManager.createBranchedSession(resumedManager.getLeafId());
  if (!forkFile) throw new Error("Prime did not create the fork session file");
  await session.disposeAsync();
  const forkManager = prime.SessionManager.open(forkFile, sessionDir, cwd);
  ({ session, extensionsResult } = await prime.createAgentSession({
    ...sessionOptions(forkManager),
    sessionStartEvent: {
      type: "session_start",
      reason: "fork",
      previousSessionFile: forkSourceFile,
    },
  }));
  await session.bindExtensions({});
  if (extensionsResult.errors.length) throw new Error(JSON.stringify(extensionsResult.errors));
  const forkEntry = prime.getLatestCompactionEntry(forkManager.getBranch());
  const forkArtifactPath = join(
    root,
    "session-artifacts",
    session.sessionId,
    "snapcompact",
    forkEntry.details.snapcompact.artifact.file,
  );
  if (!existsSync(forkArtifactPath)) {
    throw new Error(`fork lifecycle did not copy the inherited artifact: ${forkArtifactPath}; source=${forkSourceFile}; fork=${forkFile}; id=${session.sessionId}`);
  }
  await session.disposeAsync();

  const child = spawnSync(
    process.execPath,
    [
      fileURLToPath(new URL("./fork-resume-child.mjs", import.meta.url)),
      forkFile,
      sessionDir,
      cwd,
      agentDir,
    ],
    {
      cwd,
      encoding: "utf8",
      env: { ...process.env, FORCE_COLOR: "0", NO_COLOR: undefined },
    },
  );
  if (child.status !== 0) {
    throw new Error(`fork restart child failed: ${child.stderr || child.stdout}`);
  }
  const forkRestart = JSON.parse(child.stdout.trim());

  console.log(JSON.stringify({
    automaticCompaction: true,
    detailsBytes,
    artifactBytes: entry.details.snapcompact.artifact.bytes,
    firstImages,
    secondImages,
    resumedImages,
    forkRestartImages: forkRestart.forkRestartImages,
  }, null, 2));
} finally {
  await session.disposeAsync();
  registration.unregister();
  rmSync(root, { recursive: true, force: true });
}