// The harness's one direct tool: model-written javascript that calls every other tool.

import type { JsonValue } from "@earendil-works/chord";
import type { ImageContent, TextContent, TSchema } from "@earendil-works/pi-ai";
import {
  type CodemodeJsonSchema,
  type CodemodeOutputItem,
  type CodemodeResult,
  CodemodeSandbox,
  CodemodeSourceError,
  parseCodemodeSource,
  renderToolOutputType,
  renderToolSignature,
  toCodemodeIdentifier,
} from "@earendil-works/pi-codemode";

export const DESCRIPTION =
  `Run raw JavaScript (not JSON, no code fence) as an async function body in a QuickJS sandbox, without Node, file system, network, or timers.
- \`await tools.<name>(args)\` calls a tool declared in the system prompt, rejecting on failure; calls still running when the script ends are cancelled.
- Only \`text(value)\`, \`image(imageBlock)\`, \`console.log(...)\`, and \`return value\` reach you; \`exit()\` ends the script.
- \`store(key, value)\` and \`load(key)\` keep JSON values across calls.
- Optional first line: \`// @options: {"max_output_tokens": 10000, "timeout_ms": 60000}\`
- Batch independent calls (\`Promise.allSettled\`) and filter large values in one script rather than many calls.`;

export type StoreWrite = { set: Record<string, JsonValue>; delete: string[] };

/** How a script's `store()` writes fold into the values `load()` reads. */
export function applyStoreWrite(
  store: Record<string, JsonValue>,
  write: StoreWrite,
): Record<string, JsonValue> {
  const next = { ...store };
  for (const key of write.delete) delete next[key];
  return Object.assign(next, write.set);
}

/** A tool as scripts see it. */
export type Declared = {
  readonly name: string;
  readonly extension: string;
  readonly description: string;
  readonly parameters: TSchema;
  readonly returns?: TSchema;
  /** Whether the model can also call it directly, and so already has its description. */
  readonly direct?: boolean;
};

export type NestedTool = Declared & {
  call(args: JsonValue, options: {
    signal: AbortSignal;
    output(chunk: string): void;
    progress(details: JsonValue): void;
  }): Promise<JsonValue | undefined>;
};

/** One nested call as the ui sees it. Values, output, and errors are clipped; the model never sees any of it. */
export type NestedCall = {
  name: string;
  extension: string;
  args: JsonValue;
  status: "running" | "ok" | "error" | "cancelled";
  /** Wall-clock milliseconds, so a reader that comes back mid-call can still tell how long it has run. */
  startedAt: number;
  durationMs?: number;
  output?: string;
  details?: JsonValue;
  value?: JsonValue;
  error?: string;
};

export type CodemodeDetails = {
  /** Wall-clock milliseconds, as a nested call's. */
  startedAt: number;
  /** Once the script has ended. */
  durationMs?: number;
  calls: NestedCall[];
  fullOutputPath?: string;
};

const MEMORY_LIMIT_BYTES = 256 * 1024 * 1024;
const DEFAULT_MAX_OUTPUT_TOKENS = 10_000;
const CHARS_PER_TOKEN = 4;
const DETAIL_CHARS = 16 * 1024;

const declaration = (tool: Declared) => ({
  name: tool.name,
  description: tool.description,
  inputSchema: tool.parameters as CodemodeJsonSchema,
  outputSchema: (tool.returns ?? { type: "string" }) as CodemodeJsonSchema,
});

/** The codemode section of the system prompt: direct tools' results and codemode-only declarations. */
export function renderCatalog(tools: readonly Declared[]): string | undefined {
  if (tools.length === 0) return undefined;
  const direct: string[] = [];
  const nested: string[] = [];
  for (const tool of tools) {
    const declared = declaration(tool);
    if (tool.direct) {
      direct.push(
        `  ${toCodemodeIdentifier(tool.name)}: ${
          renderToolOutputType(declared.outputSchema).replaceAll("\n", "\n  ")
        };`,
      );
    } else {
      const doc = tool.description.trim().replaceAll("*/", "*\\/").replaceAll("\n", "\n   * ");
      nested.push(
        `${doc ? `  /** ${doc} */\n` : ""}  ${
          renderToolSignature(declared).replaceAll("\n", "\n  ")
        }`,
      );
    }
  }
  const blocks = [
    ...(direct.length ? [`type DirectToolResults = {\n${direct.join("\n")}\n};`] : []),
    ...(nested.length ? [`declare const tools: {\n${nested.join("\n")}\n};`] : []),
  ];
  return "Codemode tools: DirectToolResults lists direct tools callable as `await tools.<name>(args)` " +
    "with their native arguments and the resolved return types. Additional tools are declared below.\n\n" +
    `\`\`\`ts\n${blocks.join("\n\n")}\n\`\`\``;
}

const clip = (
  text: string,
) => (text.length > DETAIL_CHARS ? `${text.slice(0, DETAIL_CHARS)}…` : text);
const messageOf = (error: unknown) => (error instanceof Error ? error.message : String(error));

/**
 * With more than one text item, each starts with a `==> text N/M <==` line, since providers join adjacent text
 * blocks; `console.*` lines follow everything else in one `<console_output>` block.
 */
function layout(output: readonly CodemodeOutputItem[]): (TextContent | ImageContent)[] {
  const total = output.filter((item) => item.type === "text" && !item.console).length;
  const items: (TextContent | ImageContent)[] = [];
  const consoleLines: string[] = [];
  let index = 0;
  for (const item of output) {
    if (item.type === "image") items.push(item);
    else if (item.console) consoleLines.push(item.text);
    else {items.push({
        type: "text",
        text: total > 1 ? `==> text ${++index}/${total} <==\n${item.text}` : item.text,
      });}
  }
  if (consoleLines.length > 0) {
    items.push({
      type: "text",
      text: `<console_output>\n${consoleLines.join("\n")}\n</console_output>`,
    });
  }
  return items;
}

/** Over budget, the text keeps its start and end, images follow it, and the whole text goes to a file. */
async function budget(
  items: (TextContent | ImageContent)[],
  maxTokens: number,
): Promise<{ items: (TextContent | ImageContent)[]; fullOutputPath?: string }> {
  const text = items.flatMap((item) => (item.type === "text" ? [item.text] : [])).join("\n");
  const chars = maxTokens * CHARS_PER_TOKEN;
  if (text.length <= chars) return { items };
  const head = text.slice(0, Math.floor(chars / 2));
  const tail = text.slice(text.length - Math.ceil(chars / 2));
  const omitted = Math.ceil((text.length - head.length - tail.length) / CHARS_PER_TOKEN);
  const fullOutputPath = await Deno.makeTempFile({ prefix: "berry-codemode-", suffix: ".txt" });
  await Deno.writeTextFile(fullOutputPath, text);
  return {
    items: [
      {
        type: "text",
        text: `${head}…${omitted} tokens truncated…${tail}\n\n[Full output: ${fullOutputPath}]`,
      },
      ...items.filter((item) => item.type === "image"),
    ],
    fullOutputPath,
  };
}

function failure(
  result: Extract<CodemodeResult, { ok: false }>,
  calls: readonly NestedCall[],
): string {
  const { error } = result;
  const head = error.kind === "script"
    ? (error.stack ?? `${error.name ?? "Error"}: ${error.message}`)
    : `Script ${
      error.kind === "sandbox"
        ? "sandbox failed"
        : error.kind === "timeout"
        ? "timed out"
        : "aborted"
    }: ${error.message}`;
  const made = calls.length === 0
    ? "No tool calls were made."
    : `Tool calls made before the failure (they are not undone): ${
      calls.map((c) => `${c.name} (${c.status})`).join(", ")
    }`;
  return `Script error:\n${head}\n\n${made}`;
}

export type ScriptResult = {
  content: (TextContent | ImageContent)[];
  isError: boolean;
  details: CodemodeDetails;
  /** Only successful scripts write. */
  storeWrites?: StoreWrite;
};

export async function runScript(source: string, { tools, store, signal, progress }: {
  tools: readonly NestedTool[];
  store: Readonly<Record<string, JsonValue>>;
  signal: AbortSignal;
  progress(details: CodemodeDetails): void;
}): Promise<ScriptResult> {
  let parsed;
  try {
    parsed = parseCodemodeSource(source);
  } catch (error) {
    if (!(error instanceof CodemodeSourceError)) throw error;
    return {
      content: [{ type: "text", text: `Script failed\n\n${error.message}` }],
      isError: true,
      details: { startedAt: Date.now(), durationMs: 0, calls: [] },
    };
  }

  const startedAt = Date.now();
  const calls: NestedCall[] = [];
  const publish = () => progress({ startedAt, calls: structuredClone(calls) });
  publish();
  const sandbox = new CodemodeSandbox({
    tools: tools.map((tool) => ({
      ...declaration(tool),
      execute: async (args, { signal }) => {
        const call: NestedCall = {
          name: tool.name,
          extension: tool.extension,
          args: args as JsonValue,
          status: "running",
          startedAt: Date.now(),
        };
        calls.push(call);
        publish();
        try {
          const value = await tool.call(args as JsonValue, {
            signal,
            output: (chunk) => {
              call.output = clip((call.output ?? "") + chunk);
              publish();
            },
            progress: (details) => {
              call.details = details;
              publish();
            },
          });
          call.status = "ok";
          if (value !== undefined && JSON.stringify(value).length <= DETAIL_CHARS) {
            call.value = value;
          }
          return value;
        } catch (error) {
          call.status = signal.aborted ? "cancelled" : "error";
          call.error = clip(messageOf(error));
          throw error;
        } finally {
          call.durationMs = Date.now() - call.startedAt;
          publish();
        }
      },
    })),
    timeoutMs: parsed.options.timeoutMs ?? Number.POSITIVE_INFINITY,
    memoryLimitBytes: MEMORY_LIMIT_BYTES,
  });

  let result: CodemodeResult;
  try {
    result = await sandbox.execute(parsed.code, { signal, store });
  } finally {
    await sandbox.close();
  }
  // still running: cut off by the script ending, a timeout, or an abort
  for (const call of calls) if (call.status === "running") call.status = "cancelled";

  const output = [...result.output];
  if (result.ok && result.value !== undefined) {
    output.push({
      type: "text",
      text: typeof result.value === "string" ? result.value : JSON.stringify(result.value),
    });
  }
  const items = layout(output);
  if (!result.ok) items.push({ type: "text", text: failure(result, calls) });
  const { items: shown, fullOutputPath } = await budget(
    items,
    parsed.options.maxOutputTokens ?? DEFAULT_MAX_OUTPUT_TOKENS,
  );

  const writes = result.ok ? result.storeWrites : undefined;
  const wrote = writes && (Object.keys(writes.set).length > 0 || writes.delete.length > 0);
  return {
    content: [{ type: "text", text: result.ok ? "Script completed" : "Script failed" }, ...shown],
    isError: !result.ok,
    details: {
      startedAt,
      durationMs: Date.now() - startedAt,
      calls,
      ...(fullOutputPath ? { fullOutputPath } : {}),
    },
    ...(wrote
      ? { storeWrites: { set: writes.set as Record<string, JsonValue>, delete: writes.delete } }
      : {}),
  };
}