import type { Plugin } from "@opencode/plugin";
import { z } from "zod";
import { getInternalBridgeConfig } from "./internalBridge.js";
import { TASK_AGENT, TASK_READER_AGENT, TASK_READER_TOOLS, TASK_SUBAGENT_LIMIT } from "../src/runtime/task-agent-policy.js";

type Context = Parameters<Plugin.Plugin["setup"]>[0];
const policySchema = z.object({ role: z.enum(["coordinator", "reader", "denied"]) });
async function policy(sessionId: string, signal?: AbortSignal) {
  const { baseUrl, internalToken } = getInternalBridgeConfig();
  const response = await fetch(`${baseUrl}/internal/operational-task/agent-policy`, {
    method: "POST", headers: { "Content-Type": "application/json", "x-agent-internal-token": internalToken },
    body: JSON.stringify({ session_id: sessionId }),
    signal: signal ? AbortSignal.any([signal, AbortSignal.timeout(5000)]) : AbortSignal.timeout(5000),
  });
  if (!response.ok) throw new Error("task agent policy unavailable");
  return policySchema.parse(await response.json()).role;
}

export async function registerTaskAgentPolicy(ctx: Context) {
  const active = new Map<string, number>();
  await ctx.tool.transform((editor) => {
    editor.update("subagent", (tool) => {
      const execute = tool.execute;
      tool.execute = async (input, context) => {
        const args = z.object({ agent: z.literal(TASK_READER_AGENT), background: z.literal(false).optional(),
          sessionID: z.never().optional() }).passthrough().parse(input);
        if (await policy(context.sessionID, context.signal) !== "coordinator") throw new Error("only an active investigation may launch subagents");
        const count = active.get(context.sessionID) ?? 0;
        if (count >= TASK_SUBAGENT_LIMIT) throw new Error("at most two investigation subagents may run concurrently");
        active.set(context.sessionID, count + 1);
        try { return await execute(args, context); }
        finally {
          const remaining = (active.get(context.sessionID) ?? 1) - 1;
          if (remaining) active.set(context.sessionID, remaining);
          else active.delete(context.sessionID);
        }
      };
    });
  });
  await ctx.permission.hook("evaluate", async (event) => {
    if (event.action !== "subagent") return;
    // Respect static denials, including administrator overrides.
    if (event.effect === "deny") return;
    try {
      if (event.agent !== TASK_AGENT || !event.resources.length
        || event.resources.some((id) => id !== TASK_READER_AGENT)
        || await policy(event.sessionID) !== "coordinator") event.effect = "deny";
    } catch { event.effect = "deny"; }
  });
  await ctx.session.hook("context", async (event) => {
    if (event.agent === TASK_READER_AGENT) {
      if (await policy(event.sessionID) !== "reader") throw new Error("task reader is not bound to an active investigation");
      for (const name of Object.keys(event.tools)) if (!(TASK_READER_TOOLS as readonly string[]).includes(name)) delete event.tools[name];
    } else if (event.agent !== TASK_AGENT || await policy(event.sessionID) !== "coordinator") {
      delete event.tools.subagent;
    }
  });
}