import { Agent } from "@opencode/schema/agent";
import { Tool } from "@opencode/schema/tool";
import { Session } from "@opencode/schema/session";
import { AbsolutePath } from "@opencode/schema/schema";
import type { Context } from "@opencode/plugin/effect/plugin";
import { Cause, Effect, Schedule } from "effect";
import { z } from "zod";
import { randomUUID } from "node:crypto";
import { Herdr, rootSchema } from "./herdr.js";
import { multiWorkGuide } from "./routing.js";

export const bindingSchema = z.object({
    owner: z.string(),
    child: z.string(),
    projectID: z.string(),
    root: rootSchema,
    directory: z.string(),
    workspaceID: z.string(),
    branch: z.string(),
    ready: z.boolean(),
});
export const bindingKey = (id: string) => `multi-work:${id}`;
const inputSchema = z
    .object({
        agent: z.string(),
        description: z.string(),
        prompt: z.string(),
        sessionID: z.string().optional(),
        background: z.boolean().optional(),
        worktree: z
            .object({ base: z.string().trim().min(1) })
            .strict()
            .optional(),
    })
    .strict();
const readTools = new Set([
    "read",
    "glob",
    "grep",
    "list",
    "todoread",
    "todowrite",
    "webfetch",
    "websearch",
    "question",
    "subagent",
]);
export const childProtocol =
    "Bound child sessions use stanp-opencode multi-work protocol 1 and must remain in their assigned worktree.";

function matches(value: string, pattern: string) {
    return new RegExp(
        `^${pattern
            .replace(/[.+^${}()|[\]\\]/g, "\\$&")
            .replace(/\*/g, ".*")
            .replace(/\?/g, ".")}$`,
        "s",
    ).test(value);
}

export const installMultiWork = (
    context: Context,
    buildPermissions: () => Agent.Info["permissions"],
    herdr = new Herdr(),
) =>
    Effect.gen(function* () {
        let eligible = false;
        const busy = new Set<string>();
        const refresh = Effect.gen(function* () {
            const root = yield* Effect.tryPromise((signal) =>
                herdr.root(context.location.directory, signal),
            ).pipe(Effect.catch(() => Effect.succeed(undefined)));
            const next = root !== undefined;
            if (next === eligible) return;
            eligible = next;
            yield* context.agent.reload();
        });
        yield* refresh;
        yield* context.agent.transform((editor) => {
            editor.remove("multi-work");
            const coordinate = editor.get("coordinate");
            if (!coordinate || !editor.get("fast-lo")) return;
            editor.update("coordinate", (agent) => {
                agent.mode = "all";
                agent.system = `${agent.system ?? ""}\n${childProtocol}`;
            });
            if (!eligible) return;
            // The base plugin captures the final registry after late configuration
            // transforms and reloads when it changes. Editor.build is too early.
            const inherited = buildPermissions();
            editor.update("multi-work", (agent) => {
                Object.assign(agent, structuredClone(coordinate));
                agent.id = Agent.ID.make("multi-work");
                agent.name = Agent.Name.make("multi-work");
                agent.mode = "primary";
                agent.description =
                    "Coordinate independent tasks in Herdr worktrees; retrieve local facts with fast-lo.";
                agent.system = multiWorkGuide;
                agent.permissions = [
                    ...coordinate.permissions.filter((rule) => rule.action !== "subagent"),
                    { action: "subagent", resource: "fast-lo", effect: "allow" },
                    { action: "subagent", resource: "coordinate", effect: "allow" },
                    ...inherited
                        .filter((rule) => matches("herdr_worktree", rule.action))
                        .map((rule) => ({ ...rule, action: "herdr_worktree" })),
                ];
            });
        });
        yield* refresh.pipe(Effect.repeat(Schedule.spaced("30 seconds")), Effect.forkScoped);
        yield* context.tool.hook("execute.before", (event) => {
            if (
                event.agent === "multi-work" &&
                !readTools.has(event.tool) &&
                !event.tool.startsWith("codegraph_")
            )
                return Effect.fail(
                    new Tool.Error({ message: "multi-work may only read, plan, and delegate" }),
                );
            return Effect.void;
        });

        const validate = (sessionID: Parameters<Context["session"]["get"]>[0]["sessionID"]) =>
            Effect.gen(function* () {
                const child = yield* context.session.get({ sessionID });
                const binding = bindingSchema.parse(
                    yield* context.storage.get(bindingKey(sessionID)),
                );
                if (
                    !binding.ready ||
                    binding.child !== sessionID ||
                    child.parentID !== binding.owner ||
                    child.agent !== "coordinate" ||
                    child.projectID !== binding.projectID ||
                    child.location.directory !== binding.directory
                )
                    throw new Error("Coordinate child does not match its ready worktree binding");
                const parent = yield* context.session.get({ sessionID: child.parentID });
                if (
                    parent.agent !== "multi-work" ||
                    parent.projectID !== binding.projectID ||
                    parent.location.directory !== binding.root.directory
                )
                    throw new Error("Worktree owner changed agent or location");
                return binding;
            });

        yield* context.session.hook("context", (event) =>
            Effect.gen(function* () {
                const session = yield* context.session.get({ sessionID: event.sessionID });
                if (event.agent === "coordinate" && session.parentID)
                    yield* validate(event.sessionID);
                if (event.agent !== "multi-work") {
                    const tool = event.tools.subagent;
                    if (tool) {
                        const properties = z
                            .record(z.string(), z.unknown())
                            .parse(tool.input.properties ?? {});
                        delete properties.worktree;
                        tool.input = { ...tool.input, properties };
                    }
                    return;
                }
                for (const name of Object.keys(event.tools))
                    if (!readTools.has(name) && !name.startsWith("codegraph_"))
                        delete event.tools[name];
                if (session.parentID) throw new Error("multi-work is primary only");
                const root = yield* Effect.tryPromise((signal) =>
                    herdr.root(session.location.directory, signal),
                );
                if (!root)
                    throw new Error("multi-work requires a live, unambiguous Herdr root checkout");
                const tool = event.tools.subagent;
                if (tool) {
                    tool.input = z.toJSONSchema(inputSchema);
                    const properties = z
                        .record(z.string(), z.unknown())
                        .parse(tool.input.properties);
                    tool.input.properties = {
                        ...properties,
                        agent: { type: "string", enum: ["fast-lo", "coordinate"] },
                    };
                    tool.description =
                        "Delegate to fast-lo locally, or coordinate in a new Herdr worktree using worktree: {base}. Continue an owned coordinate child with sessionID and no worktree argument. Native background and results apply.";
                }
            }).pipe(Effect.orDie),
        );

        yield* context.tool.transform((editor) => {
            const native = editor.get("subagent");
            if (!native) return;
            editor.update("subagent", (tool) => {
                tool.input = inputSchema;
                tool.execute = (raw, call) =>
                    Effect.gen(function* () {
                        const input = inputSchema.parse(raw);
                        if (input.agent === "multi-work")
                            throw new Error("multi-work is primary only");
                        if (
                            input.sessionID &&
                            (yield* context.storage.get(bindingKey(input.sessionID))) !==
                                undefined &&
                            (call.agent !== "multi-work" || input.agent !== "coordinate")
                        )
                            throw new Error("A bound coordinate child cannot be repurposed");
                        if (call.agent !== "multi-work") {
                            if (input.worktree || input.agent === "coordinate")
                                throw new Error(
                                    "Only multi-work may initialize coordinate worktrees",
                                );
                            return yield* native.execute(input, call);
                        }
                        const parent = yield* context.session.get({ sessionID: call.sessionID });
                        if (parent.parentID || parent.agent !== "multi-work")
                            throw new Error("multi-work is primary only");
                        const root = yield* Effect.tryPromise((signal) =>
                            herdr.root(parent.location.directory, signal),
                        );
                        if (!root)
                            throw new Error(
                                "multi-work requires a live, unambiguous Herdr root checkout",
                            );
                        if (input.agent === "fast-lo") {
                            if (input.worktree) throw new Error("fast-lo is local retrieval only");
                            if (input.sessionID) {
                                const child = yield* context.session.get({
                                    sessionID: Session.ID.make(input.sessionID),
                                });
                                if (
                                    child.parentID !== parent.id ||
                                    child.agent !== "fast-lo" ||
                                    child.location.directory !== parent.location.directory
                                )
                                    throw new Error("Local fast-lo child binding changed");
                            }
                            return yield* native.execute(input, call);
                        }
                        if (input.agent !== "coordinate")
                            throw new Error("multi-work delegates only to fast-lo or coordinate");
                        if (
                            input.sessionID
                                ? input.worktree !== undefined
                                : input.worktree === undefined
                        )
                            throw new Error(
                                "New coordinate work requires worktree: {base}; resumes require only sessionID",
                            );
                        const { data: agent } = yield* context.agent.get({
                            agentID: call.agent,
                            location: { directory: parent.location.directory },
                        });
                        const permission = [
                            ...agent.permissions,
                            ...(parent.permissions ?? []),
                        ].findLast(
                            (rule) =>
                                matches("herdr_worktree", rule.action) &&
                                matches(root.directory, rule.resource),
                        )?.effect;
                        if (permission !== "allow")
                            throw new Error(
                                `herdr_worktree requires explicit allow; custom approval is unavailable (${permission ?? "ask"})`,
                            );
                        const lock = input.sessionID ?? `${call.sessionID}:${call.id}`;
                        if (busy.has(lock)) throw new Error("This delegation is already active");
                        busy.add(lock);
                        let childLock: string | undefined;
                        return yield* Effect.gen(function* () {
                            let initialized = false;
                            let placement: z.infer<typeof bindingSchema> | undefined;
                            if (input.sessionID) {
                                const binding = yield* validate(Session.ID.make(input.sessionID));
                                if (binding.owner !== call.sessionID)
                                    throw new Error("Cannot resume another parent's worktree");
                                yield* Effect.tryPromise((signal) =>
                                    herdr.verify(
                                        binding.directory,
                                        binding.workspaceID,
                                        binding.root.repo,
                                        binding.branch,
                                        signal,
                                    ),
                                );
                                placement = binding;
                                initialized = true;
                            }
                            const { worktree, ...args } = input;
                            const callKey = `multi-work-call:${call.sessionID}:${call.id}`;
                            if (!input.sessionID) {
                                const previous = yield* context.storage.get(callKey);
                                if (previous !== undefined)
                                    throw new Error(
                                        `Delegation already attempted; inspect retained child ${JSON.stringify(previous)} before retrying`,
                                    );
                            }
                            const result = yield* native.execute(args, {
                                ...call,
                                progress: (metadata) =>
                                    Effect.gen(function* () {
                                        // Native 2.0.4 awaits this update after creation and before prompt admission.
                                        yield* call.progress(metadata);
                                        if (initialized) return;
                                        const id = z.string().parse(metadata.sessionID);
                                        if (busy.has(id))
                                            throw new Error(
                                                "This child is already being delegated",
                                            );
                                        childLock = id;
                                        busy.add(id);
                                        if (!worktree)
                                            throw new Error("Missing worktree initialization");
                                        const sessionID = Session.ID.make(id);
                                        const child = yield* context.session.get({ sessionID });
                                        if (
                                            child.parentID !== call.sessionID ||
                                            child.agent !== "coordinate" ||
                                            child.projectID !== parent.projectID ||
                                            child.location.directory !== root.directory
                                        )
                                            throw new Error(
                                                "Native executor returned an unexpected child",
                                            );
                                        const branch = `multi-work/${randomUUID()}`;
                                        yield* context.storage.set(callKey, {
                                            child: id,
                                            branch,
                                            root,
                                        });
                                        yield* call.progress({ ...metadata, branch });
                                        const destination = yield* Effect.tryPromise((signal) =>
                                            herdr.create(root, worktree.base, branch, signal),
                                        );
                                        const binding = {
                                            owner: call.sessionID,
                                            child: id,
                                            projectID: parent.projectID,
                                            root,
                                            ...destination,
                                            ready: false,
                                        };
                                        yield* context.storage.set(bindingKey(id), binding);
                                        yield* call.progress({ ...metadata, ...destination });
                                        yield* Effect.gen(function* () {
                                            // Cold locations activate plugins asynchronously.
                                            yield* Effect.gen(function* () {
                                                for (;;) {
                                                    const { data: agents } =
                                                        yield* context.agent.list({
                                                            location: {
                                                                directory: destination.directory,
                                                            },
                                                        });
                                                    const coordinator = agents.find(
                                                        (agent) => agent.id === "coordinate",
                                                    );
                                                    if (
                                                        coordinator?.mode === "all" &&
                                                        coordinator.system?.includes(
                                                            childProtocol,
                                                        ) &&
                                                        [
                                                            "coordinate",
                                                            "fast-lo",
                                                            "fast-hi",
                                                            "good-lo",
                                                            "good-hi",
                                                        ].every((name) =>
                                                            agents.some(
                                                                (agent) => agent.id === name,
                                                            ),
                                                        )
                                                    )
                                                        return;
                                                    yield* Effect.sleep("100 millis");
                                                }
                                            }).pipe(Effect.timeout("30 seconds"));
                                            yield* context.session.move({
                                                sessionID,
                                                directory: AbsolutePath.make(destination.directory),
                                            });
                                            yield* Effect.gen(function* () {
                                                while (
                                                    (yield* context.session.get({ sessionID }))
                                                        .location.directory !==
                                                    destination.directory
                                                )
                                                    yield* Effect.sleep("50 millis");
                                            }).pipe(Effect.timeout("30 seconds"));
                                            if (
                                                (yield* context.session.get({ sessionID }))
                                                    .projectID !== parent.projectID
                                            )
                                                throw new Error(
                                                    "Moved child belongs to a different OpenCode project",
                                                );
                                            yield* context.storage.set(bindingKey(id), {
                                                ...binding,
                                                ready: true,
                                            });
                                            placement = { ...binding, ready: true };
                                            initialized = true;
                                        }).pipe(
                                            Effect.catchCause((cause) => {
                                                if (Cause.hasInterrupts(cause))
                                                    return Effect.interrupt;
                                                return Effect.fail(
                                                    new Tool.Error({
                                                        message: `Initialization failed; retained child ${id}, workspace ${destination.workspaceID}, checkout ${destination.directory}`,
                                                        error: cause,
                                                    }),
                                                );
                                            }),
                                        );
                                    }).pipe(Effect.orDie),
                            });
                            if (!initialized || !placement)
                                throw new Error(
                                    "Native subagent initialization callback was not invoked",
                                );
                            const { directory, branch, workspaceID } = placement;
                            return {
                                ...result,
                                metadata: { ...result.metadata, directory, branch, workspaceID },
                                content: `${z.string().parse(result.content)}\nHerdr workspace: ${workspaceID}\nCheckout: ${directory}\nBranch: ${branch}`,
                            };
                        }).pipe(
                            Effect.ensuring(
                                Effect.sync(() => {
                                    busy.delete(lock);
                                    if (childLock) busy.delete(childLock);
                                }),
                            ),
                        );
                    }).pipe(
                        Effect.catchCause((cause) => {
                            // Interruption must remain interruption so native cancellation finalizers run.
                            if (Cause.hasInterrupts(cause)) return Effect.interrupt;
                            return Effect.fail(
                                new Tool.Error({ message: Cause.pretty(cause), error: cause }),
                            );
                        }),
                    );
            });
        });
    });