import { Agent, Model, Plugin, Provider } from "@opencode/plugin";
import type { DeepMutable } from "@opencode/plugin/promise/types";
import type { ModelInfo } from "@opencode/client";
import { z } from "zod";
import { settings, workerNames } from "./config.js";
import { model } from "./model.js";
import { discover, modelKey } from "./provider.js";
import { descriptions, routingGuide } from "./routing.js";
import { Skill } from "@opencode/schema/skill";
import { AbsolutePath } from "@opencode/schema/schema";
import antiOverengineering from "../skills/anti-overengineering.md" with { type: "text" };
import taskPlanning from "../skills/task-planning.md" with { type: "text" };
import { AgentsRpc } from "./agents-rpc.js";
import { matches } from "./wildcard.js";

// Assignments use "" as unset.
function assigned(value: string | undefined): value is string {
    return value !== undefined && value !== "";
}

/** Points the agent at the available model, using the variant or a reasoning-effort fallback. */
function useModel(
    agent: DeepMutable<Agent.Info>,
    assignment: ReturnType<typeof model>,
    available: ModelInfo,
) {
    const hasVariant =
        assignment.variant !== undefined &&
        available.variants.some(({ id }) => id === assignment.variant);
    agent.model = {
        providerID: Provider.ID.make(available.providerID),
        id: Model.ID.make(available.id),
    };
    if (hasVariant) agent.model.variant = Model.VariantID.make(assignment.variant);
    agent.request = {
        settings:
            hasVariant || assignment.variant === undefined
                ? {}
                : { reasoningEffort: assignment.variant },
        headers: {},
        body: {},
    };
}

const reading = [
    "read",
    "glob",
    "grep",
    "list",
    "todoread",
    "todowrite",
    "webfetch",
    "websearch",
    "question",
    "skill",
    "external_directory",
    "doom_loop",
];
const delegation = z.object({ agent: z.string() });
export default Plugin.define({
    id: "stanp-opencode.agents",
    async setup(context) {
        const { agents } = settings.parse(context.options);
        if (!workerNames.some((name) => assigned(agents?.[name]))) return;
        const assignments: Partial<Record<(typeof workerNames)[number], ReturnType<typeof model>>> =
            {};
        for (const name of workerNames) {
            const reference = agents?.[name];
            if (assigned(reference)) assignments[name] = model(reference, `agents ${name}`);
            else if (!name.startsWith("deep-")) throw new Error(`agents ${name} is required`);
        }
        let models = new Map<string, ModelInfo>();
        let inherited: Agent.Info["permissions"] = Agent.Info.default(
            Agent.ID.make("build"),
        ).permissions;
        let permissionKey = JSON.stringify(inherited);
        let codegraphTools: string[] = [];
        let agentTransformReady = false;
        const controller = new AbortController();
        // Subscribe before the first read: configuration can publish its models while
        // asynchronous setup is running.
        const updates = (async () => {
            for await (const event of context.event.subscribe({ signal: controller.signal })) {
                if (event.location) {
                    const workspaceID =
                        "workspaceID" in event.location ? event.location.workspaceID : undefined;
                    if (
                        event.location.directory !== context.location.directory ||
                        workspaceID !== context.location.workspaceID
                    )
                        continue;
                }
                if (event.type === "model.updated") {
                    models = await discover(context);
                    await context.agent.reload();
                } else if (event.type === "agent.updated") {
                    // The configured build transform may run after ours. Read the final
                    // registry, then replay only when its effective permissions change.
                    const { data } = await context.agent.list();
                    const permissions =
                        data.find((agent) => agent.id === "build")?.permissions ??
                        Agent.Info.default(Agent.ID.make("build")).permissions;
                    const key = JSON.stringify(permissions);
                    if (key === permissionKey) continue;
                    inherited = permissions;
                    permissionKey = key;
                    await context.agent.reload();
                }
            }
        })().catch((error) => {
            if (!controller.signal.aborted)
                console.error("stanp-opencode model update failed", error);
        });
        models = await discover(context);
        const { data: initialAgents } = await context.agent.list();
        inherited = initialAgents.find((agent) => agent.id === "build")?.permissions ?? inherited;
        permissionKey = JSON.stringify(inherited);
        const registration = await context.rpc.register(AgentsRpc, {
            permissions: async () => inherited.map((rule) => ({ ...rule })),
        });
        let workers: (typeof workerNames)[number][] = [];
        const active = new Set<string>();

        await context.skill.transform((editor) => {
            for (const [name, content] of [
                ["anti-overengineering", antiOverengineering],
                ["task-planning", taskPlanning],
            ] as const) {
                if (editor.get(name)) continue;
                const description = content.match(/^description: "?(.+?)"?$/mu)?.[1];
                if (description === undefined || description === "")
                    throw new Error(`${name} skill has no description`);
                editor.add({
                    id: Skill.ID.make(name),
                    name: Skill.Name.make(name),
                    description,
                    path: AbsolutePath.make(import.meta.filename),
                    content,
                });
            }
        });

        await context.tool.transform((editor) => {
            const next = editor
                .list()
                .map((tool) => tool.id)
                .filter((id) => id.startsWith("codegraph_"))
                .toSorted();
            if (JSON.stringify(next) === JSON.stringify(codegraphTools)) return;
            codegraphTools = next;
            if (agentTransformReady)
                queueMicrotask(() => {
                    void context.agent.reload().catch((error) => {
                        console.error("stanp-opencode agent reload failed", error);
                    });
                });
        });

        const availableFor = (assignment: { providerID: string; id: string }) =>
            models.get(modelKey(assignment.providerID, assignment.id));
        await context.agent.transform((editor) => {
            const deepLo = assignments["deep-lo"];
            const deep =
                deepLo !== undefined &&
                assignments["deep-hi"] !== undefined &&
                availableFor(deepLo) !== undefined;
            const candidates = workerNames.filter((name) => deep || !name.startsWith("deep-"));
            const coordinate = deep ? deepLo : assignments["good-lo"];
            workers = [coordinate, ...candidates.map((name) => assignments[name])].every(
                (assignment) => assignment !== undefined && availableFor(assignment) !== undefined,
            )
                ? candidates
                : [];
            active.clear();
            for (const name of workers) active.add(name);
            if (workers.length === 0 || coordinate === undefined) return;
            const configureWorker = (
                name: (typeof workerNames)[number],
                assignment: ReturnType<typeof model>,
                available: ModelInfo,
            ) => {
                editor.update(name, (agent) => {
                    agent.name = Agent.Name.make(name);
                    agent.mode = "subagent";
                    agent.hidden = false;
                    agent.description = descriptions[name];
                    useModel(agent, assignment, available);
                    agent.permissions = [
                        ...structuredClone(inherited),
                        { action: "subagent", resource: "*", effect: "deny" },
                    ];
                });
            };
            for (const name of workers) {
                const assignment = assignments[name];
                const available = assignment && availableFor(assignment);
                if (assignment && available) configureWorker(name, assignment, available);
            }
            const available = availableFor(coordinate);
            if (!available) return;
            editor.update("coordinate", (agent) => {
                agent.name = Agent.Name.make("coordinate");
                agent.mode = "primary";
                agent.hidden = false;
                agent.description =
                    "Read, plan, and coordinate specialist workers; never implement directly.";
                useModel(agent, coordinate, available);
                agent.system = routingGuide(deep);
                agent.permissions = [
                    { action: "*", resource: "*", effect: "deny" },
                    ...reading.flatMap((action) =>
                        inherited
                            .filter((rule) => matches(action, rule.action))
                            .map((rule) => ({ ...rule, action })),
                    ),
                    ...codegraphTools.flatMap((action) =>
                        inherited
                            .filter((rule) => matches(action, rule.action))
                            .map((rule) => ({ ...rule, action })),
                    ),
                    ...workers.map((resource) => ({
                        action: "subagent",
                        resource,
                        effect: "allow" as const,
                    })),
                ];
            });
        });
        agentTransformReady = true;
        async function assertMultiWork() {
            const { data } = await context.plugin.list();
            if (
                !data.some(
                    (plugin) =>
                        plugin.id === "stanp-opencode.multi-work" &&
                        plugin.state.status === "active",
                )
            )
                throw new Error("coordinate cannot run as a subagent without Multi-work");
        }
        await context.session.hook("context", (event) => {
            if (workerNames.some((name) => name === event.agent)) {
                delete event.tools["subagent"];
                return;
            }
            const tool = event.tools["subagent"];
            if (event.agent === "coordinate" && workers.length === 0) {
                delete event.tools["subagent"];
                return;
            }
            if (event.agent !== "coordinate" || !tool) return;
            tool.description = `Delegate a self-contained assignment to one of these workers:\n${workers.map((name) => `- ${name}: ${descriptions[name]}`).join("\n")}`;
            // The executable definition remains native; constrain only this request's schema.
            const properties = z
                .record(z.string(), z.unknown())
                .parse(tool.input["properties"] ?? {});
            tool.input = {
                ...tool.input,
                properties: { ...properties, agent: { type: "string", enum: workers } },
            };
        });
        await context.tool.hook("execute.before", async (event) => {
            if (event.tool !== "subagent") return;
            const { agent } = delegation.parse(event.input);
            if (agent === "coordinate" && event.agent !== "multi-work")
                throw new Error("coordinate cannot run as a subagent without Multi-work");
            if (agent === "coordinate") await assertMultiWork();
            if (workerNames.some((name) => name === event.agent))
                throw new Error("Workers cannot delegate tasks");
            if (event.agent === "coordinate" && !active.has(agent))
                throw new Error("coordinate may delegate only to its available workers");
        });
        // Commands can create child sessions without going through the subagent tool.
        await context.session.hook("prompt", async (event) => {
            const child = await context.session.get({ sessionID: event.sessionID });
            if (child.parentID === undefined || child.parentID === "") return;
            const parent = await context.session.get({ sessionID: child.parentID });
            if (child.agent === "coordinate" && parent.agent !== "multi-work")
                throw new Error("coordinate cannot run as a subagent without Multi-work");
            if (child.agent === "coordinate") await assertMultiWork();
            if (workerNames.some((name) => name === parent.agent))
                throw new Error("Workers cannot delegate tasks");
            if (
                parent.agent === "coordinate" &&
                (child.agent === undefined || !active.has(child.agent))
            )
                throw new Error("coordinate may delegate only to its available workers");
        });
        return async () => {
            controller.abort();
            await updates;
            await registration.dispose();
        };
    },
});