import type { Sandbox } from "@cloudflare/sandbox";
import { SandboxAgent } from "sandbox-agent";

export type PromptRequest = {
  agent?: string;
  prompt?: string;
};

export async function runPromptEndpointStream(
  sandbox: Sandbox,
  request: PromptRequest,
  port: number,
  emit: (event: { type: string; [key: string]: unknown }) => Promise<void> | void,
): Promise<void> {
  const client = await SandboxAgent.connect({
    fetch: (req, init) =>
      sandbox.containerFetch(
        req,
        {
          ...(init ?? {}),
          // Cloudflare containerFetch may drop long-lived update streams when
          // a forwarded AbortSignal is cancelled; clear it for this path.
          signal: undefined,
        },
        port,
      ),
  });

  let unsubscribe: (() => void) | undefined;
  try {
    const session = await client.createSession({
      agent: request.agent ?? "codex",
    });

    const promptText = request.prompt?.trim() || "Reply with a short confirmation.";
    await emit({
      type: "session.created",
      sessionId: session.id,
      agent: session.agent,
      prompt: promptText,
    });

    let pendingWrites: Promise<void> = Promise.resolve();
    unsubscribe = session.onEvent((event) => {
      pendingWrites = pendingWrites
        .then(async () => {
          await emit({ type: "session.event", event });
        })
        .catch(() => {});
    });

    const response = await session.prompt([{ type: "text", text: promptText }]);
    await pendingWrites;
    await emit({ type: "prompt.response", response });
    await emit({ type: "prompt.completed" });
  } finally {
    if (unsubscribe) {
      unsubscribe();
    }
    await Promise.race([client.dispose(), new Promise((resolve) => setTimeout(resolve, 250))]);
  }
}