import { afterEach, describe, expect, it } from "bun:test";
import { registerTaskAgentPolicy } from "../../.opencode/task-agent-policy.js";
import { TASK_AGENT, TASK_READER_AGENT } from "../../src/runtime/task-agent-policy.js";

const originalFetch = globalThis.fetch;
const originalToken = process.env.TJWATER_AGENT_INTERNAL_TOKEN;
afterEach(() => {
  globalThis.fetch = originalFetch;
  if (originalToken === undefined) delete process.env.TJWATER_AGENT_INTERNAL_TOKEN;
  else process.env.TJWATER_AGENT_INTERNAL_TOKEN = originalToken;
});

describe("investigation subagent plugin", () => {
  it("bounds parallel calls, rejects background/resume/other agents, and releases slots on failure", async () => {
    process.env.TJWATER_AGENT_INTERNAL_TOKEN = "TEST_ONLY";
    let role = "coordinator";
    globalThis.fetch = (async () => Response.json({ role })) as unknown as typeof fetch;
    const pending: Array<() => void> = [];
    let entered = 0;
    const tool = { execute: async () => {
      entered++;
      await new Promise<void>((resolve) => pending.push(resolve));
      throw new Error("native failure");
    } };
    const hooks = new Map<string, (event: any) => Promise<void>>();
    await registerTaskAgentPolicy({
      tool: { transform: async (fn: any) => fn({ update: (_name: string, update: any) => update(tool) }) },
      permission: { hook: async (_name: string, fn: any) => { hooks.set("permission", fn); } },
      session: { hook: async (_name: string, fn: any) => { hooks.set("context", fn); } },
    } as never);
    const execute = tool.execute as unknown as (input: unknown, context: unknown) => Promise<unknown>;
    const call = (extra = {}) => execute({ agent: TASK_READER_AGENT, prompt: "check", description: "check", ...extra }, { sessionID: "parent", signal: new AbortController().signal });
    await expect(call({ background: true })).rejects.toThrow();
    await expect(call({ sessionID: "foreign-child" })).rejects.toThrow();
    await expect(call({ agent: "general" })).rejects.toThrow();
    const first = call().catch((error) => error);
    const second = call().catch((error) => error);
    await new Promise((resolve) => setTimeout(resolve, 0));
    expect(entered).toBe(2);
    await expect(call()).rejects.toThrow("at most two");
    pending.splice(0).forEach((finish) => finish());
    await Promise.all([first, second]);
    const third = call().catch((error) => error);
    await new Promise((resolve) => setTimeout(resolve, 0));
    expect(entered).toBe(3);
    pending.splice(0).forEach((finish) => finish());
    await third;
    role = "denied";
    await expect(call()).rejects.toThrow("only an active investigation");
    const permission = { sessionID: "ordinary", agent: "instruction", action: "subagent", resources: [TASK_READER_AGENT], effect: "allow" };
    await hooks.get("permission")!(permission);
    expect(permission.effect).toBe("deny");
    role = "coordinator";
    Object.assign(permission, { agent: TASK_AGENT, effect: "deny" });
    await hooks.get("permission")!(permission);
    expect(permission.effect).toBe("deny");
    role = "reader";
    const event = { sessionID: "child", agent: TASK_READER_AGENT, tools: {
      mcp_tjwater_get_task_context: {}, mcp_tjwater_query_scada_readings: {}, publish_task_report: {}, subagent: {}, execute: {}, shell: {},
    } };
    await hooks.get("context")!(event);
    expect(Object.keys(event.tools)).toEqual(["mcp_tjwater_get_task_context", "mcp_tjwater_query_scada_readings"]);
  });
});