import runtime, {
  PersonalAgent as FixturePersonalAgent,
  Conversation,
} from "./think-worker";
import { getAgentByName } from "agents";
import type { Env } from "../../worker/personal-agent";
import { Sandbox as RuntimeSandbox } from "../../worker/sandbox";
export { Conversation };

let launchDelay = 0;
let nativeLaunchDelay = 0;
let destroyFailures = 0;
const observed: string[] = [];
const events: { id: string; event: string }[] = [];

// Delay inside the native Sandbox RPC lifetime, after the one-use lease gate.
export class Sandbox extends RuntimeSandbox {
  async execStreamWithSessionToken(
    ...args: Parameters<RuntimeSandbox["execStreamWithSessionToken"]>
  ) {
    const delay = nativeLaunchDelay;
    nativeLaunchDelay = 0;
    if (delay) {
      events.push({
        id: this.ctx.id.toString(),
        event: "native-launch-waiting",
      });
      await new Promise((resolve) => setTimeout(resolve, delay));
    }
    return super.execStreamWithSessionToken(...args);
  }
}

export class PersonalAgent extends FixturePersonalAgent {
  protected shellSandbox(id: string) {
    const sandbox = super.shellSandbox(id);
    return new Proxy(sandbox, {
      get(target, key) {
        if (key === "execTemporary")
          return async (...args: Parameters<typeof sandbox.execTemporary>) => {
            events.push({ id, event: "launch-start" });
            const delay = launchDelay;
            launchDelay = 0;
            if (delay)
              await new Promise((resolve) => setTimeout(resolve, delay));
            const stream = await target.execTemporary(...args);
            events.push({ id, event: "launch-settled" });
            return stream;
          };
        if (key === "closeTemporary")
          return async () => {
            events.push({ id, event: "destroy-start" });
            if (destroyFailures-- > 0)
              throw new Error("Fixture cleanup failure");
            const closed = await target.closeTemporary();
            events.push({ id, event: "destroy-settled" });
            return closed;
          };
        return Reflect.get(target, key);
      },
    });
  }
  async reserveShellWorkspace(
    ...args: Parameters<FixturePersonalAgent["reserveShellWorkspace"]>
  ) {
    const lease = await super.reserveShellWorkspace(...args);
    observed.push(lease.id);
    return lease;
  }
  restartShellParent() {
    this.ctx.abort("Fixture parent restart");
  }
  async inspectShells() {
    return {
      leases: this.sql`SELECT * FROM flarebot_shell_leases`,
      containers: await Promise.all(
        observed.map(async (id) => ({
          id,
          state: await super.shellSandbox(id).getState(),
        })),
      ),
      events,
    };
  }
}

export default {
  async fetch(request: Request, env: Env, ctx: ExecutionContext) {
    const url = new URL(request.url);
    if (url.pathname === "/__fixture/shell-status") {
      const parent = await getAgentByName(env.PersonalAgent, "personal");
      return Response.json(
        await (parent as unknown as PersonalAgent).inspectShells(),
      );
    }
    if (url.pathname === "/__fixture/shell-fault") {
      launchDelay = Number(url.searchParams.get("launch") ?? 0);
      nativeLaunchDelay = Number(url.searchParams.get("nativeLaunch") ?? 0);
      destroyFailures = Number(url.searchParams.get("destroy") ?? 0);
      return new Response("ok");
    }
    if (url.pathname === "/__fixture/shell-restart") {
      const parent = await getAgentByName(env.PersonalAgent, "personal");
      await (parent as unknown as PersonalAgent)
        .restartShellParent()
        .catch(() => {});
      return new Response("restarted");
    }
    return runtime.fetch(request, env, ctx);
  },
};