import type { AgentRunRequest } from "@rakazo/adapter-kit";
import { describe, expect, it, vi } from "vitest";
import type * as ComputerLifecycleModule from "./computer-lifecycle.js";
import { createRunExecutor, MAIN_ONLY_BUILTIN_TOOLS } from "./executor.js";

vi.mock("./computer-lifecycle.js", async (importOriginal) => ({
  ...(await importOriginal<typeof ComputerLifecycleModule>()),
  acquireComputerExecutionLease: async () => null,
  provisionComputer: async () => ({ id: "computer-1", kind: "desktop" }),
}));

type BotPolicy = { kind?: string; restrictBuiltinTools?: boolean; allowedBuiltinTools?: string[] };
type Call = { name: string; args?: Record<string, unknown> };
type Effect = { id: string; status: string; idempotencyKey: string } & Record<string, unknown>;

/** In-memory external effects, enough for the idempotency record of a mutating builtin tool. */
function effectStore() {
  const effects: Effect[] = [];
  const find = (where: { id?: string; idempotencyKey?: string }) =>
    effects.find((effect) =>
      where.id ? effect.id === where.id : effect.idempotencyKey === where.idempotencyKey,
    );
  return {
    findMany: vi.fn(async () => []),
    findUnique: vi.fn(async ({ where }: { where: { id?: string; idempotencyKey?: string } }) => {
      const effect = find(where);
      return effect ? { ...effect } : null;
    }),
    create: vi.fn(async ({ data }: { data: Omit<Effect, "id"> }) => {
      const effect = { ...data, id: `effect-${effects.length + 1}` } as Effect;
      effects.push(effect);
      return { ...effect };
    }),
    update: vi.fn(async ({ where, data }: { where: { id: string }; data: Partial<Effect> }) => {
      Object.assign(find(where) ?? {}, data);
    }),
    updateMany: vi.fn(
      async ({ where, data }: { where: { id: string; status: string }; data: Partial<Effect> }) => {
        const effect = find(where);
        if (!effect || effect.status !== where.status) return { count: 0 };
        Object.assign(effect, data);
        return { count: 1 };
      },
    ),
  };
}

async function runBot(
  policy: BotPolicy,
  calls: Call[] = [],
  models: Record<string, unknown> = {},
  deps: Record<string, unknown> = {},
) {
  const run = {
    id: "run-1",
    botId: "bot-1",
    threadId: "thread-1",
    taskId: "task-1",
    spaceId: "space-1",
    userId: "user-1",
    status: "queued",
    trigger: "user",
    leaseFence: 0,
  };
  const routineFindMany = vi.fn(async () => []);
  const prisma = {
    run: {
      findUnique: vi.fn(async () => run),
      findUniqueOrThrow: vi.fn(async () => run),
      updateMany: vi.fn(async ({ data }: { data: Record<string, unknown> }) => {
        Object.assign(run, data);
        return { count: 1 };
      }),
    },
    bot: {
      findUniqueOrThrow: vi.fn(async () => ({
        id: run.botId,
        name: "Assistant",
        title: "Assistant",
        description: "Test assistant",
        computerId: "computer-1",
        computer: { id: "computer-1", scope: "dedicated" },
        kind: "specialist",
        restrictBuiltinTools: false,
        allowedBuiltinTools: [],
        ...policy,
      })),
      findMany: vi.fn(async () => []),
    },
    attempt: {
      create: vi.fn(async () => ({ id: "attempt-1" })),
      update: vi.fn(),
      updateMany: vi.fn(async () => ({ count: 1 })),
    },
    thread: { findUniqueOrThrow: vi.fn(async () => ({ id: run.threadId, groupId: null })) },
    space: { findUnique: vi.fn(async () => ({ name: "Home", isDefault: true })) },
    message: { findMany: vi.fn(async () => []) },
    task: { findUniqueOrThrow: vi.fn(async () => ({ id: run.taskId, prompt: "Work" })) },
    connection: { findMany: vi.fn(async () => []) },
    spaceModelPreference: { findFirst: vi.fn(async () => null) },
    userModelCredential: { findFirst: vi.fn(async () => null) },
    deploymentSettings: {
      findUnique: vi.fn(async () => ({
        defaultModelProvider: "scripted",
        defaultModelId: "scripted",
      })),
    },
    taughtSkill: { findMany: vi.fn(async () => []) },
    agentSecret: { findMany: vi.fn(async () => []) },
    agentSkill: { findMany: vi.fn(async () => []) },
    scratchpadItem: { findMany: vi.fn(async () => []) },
    actionApprovalRule: { findMany: vi.fn(async () => []) },
    actionAutoReviewPreference: { findUnique: vi.fn(async () => ({ enabled: false })) },
    externalEffect: effectStore(),
    routine: { findMany: routineFindMany },
    // The user block every run states.
    user: { findUnique: vi.fn(async () => ({ name: "Ada" })) },
    userProfile: { findUnique: vi.fn(async () => null) },
    userFact: { findMany: vi.fn(async () => []) },
    // The main bot's goals and habits.
    userGoal: { findMany: vi.fn(async () => []) },
    habit: { findMany: vi.fn(async () => []) },
    ...models,
  };
  let request: AgentRunRequest | undefined;
  const results: unknown[] = [];
  const jobs = { enqueue: vi.fn(async () => undefined) };
  const executor = createRunExecutor({
    prisma,
    runtime: {
      describe: () => ({ capabilities: { scripted: false } }),
      run: async function* (next: AgentRunRequest) {
        request = next;
        for (const [index, call] of calls.entries()) {
          results.push(await next.executeTool!(call.name, call.args ?? {}, `call-${index}`));
        }
        yield { type: "done" as const, text: "Done" };
      },
    },
    sandbox: { describe: () => ({ capabilities: { graphical: false } }) },
    memory: { read: async () => ({ documents: [] }) },
    memoryProviders: { resolve: async () => null },
    events: {
      append: vi.fn(async () => undefined),
      pauseRunForInput: vi.fn(async () => true),
      finalizeRun: vi.fn(async () => ({ continuationRunId: null })),
    },
    jobs,
    secrets: [],
    ...deps,
  } as unknown as Parameters<typeof createRunExecutor>[0]);
  await executor.continueRun(run.id, "worker-1");
  if (!request) throw new Error("the runtime was never started");
  const tools = request.tools;
  return { names: tools.map((tool) => tool.name), tools, results, routineFindMany };
}

describe("bot capabilities in runs", () => {
  it("exposes every builtin to an unrestricted bot", async () => {
    const { names } = await runBot({});
    expect(names).toEqual(expect.arrayContaining(["write_file", "shell", "schedule_list"]));
    // Media tools need a media runtime that can run their operation.
    expect(names).not.toContain("generate_image");
  });

  it("exposes media tools the space can run and starts their jobs without waiting", async () => {
    const start = vi.fn(async () => ({ jobId: "job-1", messageId: "message-1" }));
    const media = {
      capabilities: async () => ({
        generate_image: true,
        generate_video: false,
        edit_image: false,
        edit_video: false,
        export_video: false,
        transcribe: false,
      }),
      start,
      get: async () => null,
      cancel: async () => undefined,
    };
    const { names, results } = await runBot(
      {},
      [
        { name: "generate_image", args: { prompt: "a red fox", continue_after: true } },
        { name: "generate_video", args: { prompt: "a red fox" } },
      ],
      {},
      { media },
    );
    expect(names).toContain("generate_image");
    expect(names).not.toContain("generate_video");
    expect(results[0]).toMatchObject({ jobId: "job-1", status: "queued" });
    expect(start).toHaveBeenCalledExactlyOnceWith({
      operation: "generate_image",
      input: { prompt: "a red fox" },
      spaceId: "space-1",
      userId: "user-1",
      botId: "bot-1",
      groupId: null,
      threadId: "thread-1",
      runId: "run-1",
      operationKey: expect.any(String),
      continueAfter: true,
    });
    expect(results[1]).toEqual({ error: "Tool generate_video is not available in this run." });
  });

  it("exposes only the allowlisted builtins and refuses calls to any other", async () => {
    const { names, results, routineFindMany } = await runBot(
      { restrictBuiltinTools: true, allowedBuiltinTools: ["schedule_list", "web_search"] },
      [
        { name: "write_file", args: { path: "notes/x.txt", content: "nope" } },
        { name: "schedule_list" },
      ],
    );
    expect(names).toEqual(["web_search", "schedule_list"]);
    expect(results[0]).toEqual({ error: "Tool write_file is not available in this run." });
    expect(results[1]).toEqual({ routines: [] });
    expect(routineFindMany).toHaveBeenCalledOnce();
  });

  it("exposes no builtins when the allowlist is empty", async () => {
    const { names, results } = await runBot({ restrictBuiltinTools: true }, [
      { name: "schedule_list" },
    ]);
    expect(names).toEqual([]);
    expect(results[0]).toEqual({ error: "Tool schedule_list is not available in this run." });
  });

  it("refuses a tool no run ever exposes", async () => {
    const { results } = await runBot({}, [{ name: "made_up_tool" }]);
    expect(results[0]).toEqual({ error: "Tool made_up_tool is not available in this run." });
  });

  it("offers schedule_create's bot option to the main bot only", async () => {
    const properties = async (policy: BotPolicy) =>
      Object.keys(
        ((await runBot(policy)).tools.find((tool) => tool.name === "schedule_create")?.inputSchema
          .properties ?? {}) as object,
      );
    expect(await properties({ kind: "main" })).toContain("bot");
    expect(await properties({})).not.toContain("bot");
  });

  it("withholds main-only tools from every other bot", async () => {
    const mainOnly = MAIN_ONLY_BUILTIN_TOOLS as Set<string>;
    // No main-only tool ships yet in this package; schedule_list stands in for one.
    mainOnly.add("schedule_list");
    try {
      const specialist = await runBot({}, [{ name: "schedule_list" }]);
      expect(specialist.names).not.toContain("schedule_list");
      expect(specialist.results[0]).toEqual({
        error: "Tool schedule_list is not available in this run.",
      });
      const main = await runBot({ kind: "main" }, [{ name: "schedule_list" }]);
      expect(main.names).toContain("schedule_list");
      expect(main.results[0]).toEqual({ routines: [] });
    } finally {
      mainOnly.delete("schedule_list");
    }
  });

  it("schedules for a named specialist when the main bot asks, in that bot's space and thread", async () => {
    const routineCreate = vi.fn(async ({ data }: { data: Record<string, unknown> }) => ({
      id: "routine-1",
      ...data,
    }));
    const { results } = await runBot(
      { kind: "main" },
      [
        {
          name: "schedule_create",
          args: { name: "Digest", prompt: "Summarize", delayMinutes: 5, bot: "Atlas" },
        },
      ],
      {
        bot: {
          findUniqueOrThrow: vi.fn(async () => ({
            id: "bot-1",
            name: "Chief",
            title: "",
            description: "",
            computerId: "computer-1",
            computer: { id: "computer-1", scope: "dedicated" },
            kind: "main",
            restrictBuiltinTools: false,
            allowedBuiltinTools: [],
          })),
          findMany: vi.fn(async () => [
            {
              id: "atlas-1",
              name: "Atlas",
              title: "",
              description: "",
              spaceId: "space-2",
              space: { name: "Studio" },
              thread: { id: "atlas-dm" },
            },
          ]),
        },
        chatGroup: { findFirst: vi.fn(async () => null) },
        routine: {
          findMany: vi.fn(async () => []),
          count: vi.fn(async () => 0),
          create: routineCreate,
        },
      },
    );
    expect(results[0]).toMatchObject({ ok: true, routineId: "routine-1", bot: "Atlas" });
    expect(routineCreate).toHaveBeenCalledWith(
      expect.objectContaining({
        data: expect.objectContaining({
          spaceId: "space-2",
          botId: "atlas-1",
          threadId: "atlas-dm",
          createdByBotId: "bot-1",
        }),
      }),
    );
  });

  it("refuses the bot option from any other bot", async () => {
    const routineCreate = vi.fn();
    const { results } = await runBot(
      {},
      [
        {
          name: "schedule_create",
          args: { name: "x", prompt: "y", delayMinutes: 5, bot: "Atlas" },
        },
      ],
      { routine: { findMany: vi.fn(async () => []), create: routineCreate } },
    );
    expect(results[0]).toEqual({ error: "Only the main agent can schedule for another bot." });
    expect(routineCreate).not.toHaveBeenCalled();
  });
});