import * as crypto from "node:crypto";
import { StringEnum } from "@earendil-works/pi-ai";
import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, truncateHead, type ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type, type Static } from "typebox";
import { ackMessage, growthProposals, messages, putGrowth, storeMessages, messageReceiptDetails, readJson, runFile, type ControlMessage, type GrowthProposal } from "./store.ts";
import type { MeshRun, MeshTask } from "./manager.ts";
import { MeshTaskSchema } from "./schemas.ts";


function boundedJson(value: unknown): string {
  const result = truncateHead(JSON.stringify(value), { maxBytes: DEFAULT_MAX_BYTES, maxLines: DEFAULT_MAX_LINES });
  return result.truncated ? `${result.content}\n[Output truncated; use narrower inbox filters or inspect the project spool.]` : result.content;
}
function boundedDetails(value: Record<string, unknown>): Record<string, unknown> {
  const bytes = Buffer.byteLength(JSON.stringify(value), "utf8");
  return bytes <= DEFAULT_MAX_BYTES ? value : { truncated: true, bytes };
}

const Params = Type.Object({
  action: StringEnum(["status", "send", "broadcast", "reply", "inbox", "ack", "grow"] as const),
  to: Type.Optional(Type.String({ maxLength: 64 })), content: Type.Optional(Type.String({ maxLength: 1_048_576 })), messageId: Type.Optional(Type.String({ maxLength: 128 })),
  reason: Type.Optional(Type.String({ maxLength: 16384 })), tasks: Type.Optional(Type.Array(MeshTaskSchema, { minItems: 1, maxItems: 16 })),
}, { additionalProperties: false });
type MeshControlParams = Static<typeof Params>;
export function createMeshControlTool(root: string, runId: string, nodeId: string, attempt: number, onMessageStored?: (message: ControlMessage) => void | Promise<void>) {
  return {
    name: "mesh_control",
    label: "Mesh Control",
    description: "Child-safe topology, mailbox, and growth proposal tool. Growth only proposes tasks; the host must approve and commit them.",
    parameters: Params,
    async execute(_id: string, rawParams: MeshControlParams): Promise<any> {
      const params = rawParams;
      const run = readJson<MeshRun>(runFile(root, runId));
      const caller = run?.nodes.find((node) => node.id === nodeId);
      if (!run || !caller || caller.status !== "running" || caller.attempt !== attempt || !["running", "paused"].includes(run.status)) throw new Error("Mesh child identity is no longer active");
      if (params.action === "status") {
        const snapshot = {
          run: { id: run.id, status: run.status, operator: run.operator, revision: run.revision, maxConcurrency: run.maxConcurrency, maxNodes: run.maxNodes },
          self: { id: caller.id, agent: caller.agent, status: caller.status, attempt: caller.attempt, dependsOn: caller.dependsOn, allowedSubagents: caller.allowedSubagents ?? [] },
          nodes: run.nodes.map((node) => ({ id: node.id, agent: node.agent, status: node.status, attempt: node.attempt, dependsOn: node.dependsOn, requestedBy: node.requestedBy })),
          pendingGrowth: growthProposals(root, runId).filter((proposal) => proposal.status === "proposed").length,
        };
        return { content: [{ type: "text", text: boundedJson(snapshot) }], details: boundedDetails(snapshot) };
      }
      if (params.action === "inbox") {
        const inbox = messages(root, runId).filter((message) => message.to === nodeId && !message.ackedAt);
        const proposals = growthProposals<MeshTask[]>(root, runId).filter((proposal) => proposal.requester === nodeId).map((proposal) => {
          const nodes = (proposal.committedNodeIds ?? []).map((id) => run.nodes.find((node) => node.id === id)).filter((node): node is MeshRun["nodes"][number] => Boolean(node));
          return { ...proposal, counts: nodes.reduce<Record<string, number>>((counts, node) => { counts[node.status] = (counts[node.status] ?? 0) + 1; return counts; }, {}),
            nodes: nodes.map((node) => ({ id: node.id, status: node.status, attempt: node.attempt, error: node.error, outputPath: node.outputPath, attemptResultPath: node.attemptResultPath, diagnosticPath: node.diagnosticPath, evidencePath: node.evidencePath })) };
        });
        return { content: [{ type: "text", text: boundedJson({ inbox, growth: proposals }) }], details: boundedDetails({ inboxCount: inbox.length, growthCount: proposals.length, inbox: inbox.slice(0, 256), growth: proposals.slice(0, 256) }) };
      }
      if (params.action === "ack") {
        if (!params.messageId || !ackMessage(root, runId, params.messageId, nodeId)) throw new Error("messageId is not an unacked message for this node");
        return { content: [{ type: "text", text: `Acknowledged ${params.messageId}.` }], details: {} };
      }
      if (run.status !== "running") throw new Error("Paused Mesh permits only status/inbox/ack during drain");
      if (params.action === "grow") {
        if (!params.reason?.trim() || !params.tasks?.length) throw new Error("reason and tasks are required for grow");
        if (caller.allowedSubagents !== "all") {
          const allowed = new Set((caller.allowedSubagents ?? []).map((name) => name.toLowerCase()));
          const denied = params.tasks.map((task) => task.agent).filter((name) => !allowed.has(name.toLowerCase()));
          if (denied.length) throw new Error(`Agent ${caller.agent} cannot request growth for: ${[...new Set(denied)].join(", ")}`);
        }
        const proposal: GrowthProposal<MeshTask[]> = {
          id: crypto.randomUUID(), runId, requester: nodeId, reason: params.reason.trim(), tasks: params.tasks as MeshTask[], status: "proposed",
          baseRevision: run.revision, requesterAttempt: attempt, createdAt: Date.now(),
        };
        if (Buffer.byteLength(JSON.stringify(proposal), "utf8") > 1024 * 1024) throw new Error("Growth proposal exceeds 1048576 bytes");
        putGrowth(root, proposal);
        return { content: [{ type: "text", text: `Growth proposed: ${proposal.id}. Host approval is required.` }], details: boundedDetails({ proposal }) };
      }

      const all = messages(root, runId);
      let recipients: string[];
      let replyTo: string | undefined;
      if (params.action === "reply") {
        if (!params.messageId) throw new Error("messageId is required for reply");
        const original = all.find((message) => message.id === params.messageId && message.to === nodeId);
        if (!original) throw new Error("Cannot reply to an unknown message");
        recipients = [original.from];
        replyTo = original.id;
      } else if (params.action === "broadcast") {
        recipients = run.nodes.map((node) => node.id).filter((id) => id !== nodeId);
      } else {
        if (!params.to) throw new Error("to is required for send");
        recipients = [params.to];
      }
      if (recipients.length > 128 || recipients.some(id => !/^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$/.test(id))) throw new Error("Invalid bounded Mesh recipients");
      if (recipients.some((id) => id !== "host" && !run.nodes.some((node) => node.id === id))) throw new Error("Unknown recipient node");
      const payloadMaxBytes = run.messagePayloadMaxBytes ?? 32 * 1024;
      const recipientUnreadMaxBytes = run.recipientUnreadMaxBytes ?? 1024 * 1024;
      if (!params.content?.trim() || Buffer.byteLength(params.content, "utf8") > payloadMaxBytes) throw new Error(`content must be 1-${payloadMaxBytes} bytes`);
      const sent: ControlMessage[] = recipients.map((to) => ({
        id: crypto.randomUUID(), runId, from: nodeId, to, content: params.content!.trim(), replyTo, source: "child", senderAttempt: attempt, createdAt: Date.now(),
      }));
      const receipt = await storeMessages(root, sent, { payloadMaxBytes, recipientUnreadMaxBytes }, onMessageStored);
      return { content: [{ type: "text", text: `Stored ${receipt.stored}/${sent.length} mailbox message(s). Inspect receipts for durability/notification warnings; storage is not business completion.${receipt.partial ? " Partial delivery." : ""} For retries, retry only not-stored recipients; inspect unknown IDs before resending.` }], details: messageReceiptDetails(receipt) };
    },
  };
}

export default function registerMeshControl(pi: ExtensionAPI): void {
  const runId = process.env.PI_MESH_RUN_ID;
  const nodeId = process.env.PI_MESH_NODE_ID;
  const attempt = Number(process.env.PI_MESH_ATTEMPT);
  const root = process.env.PI_MESH_ROOT;
  if (runId && nodeId && Number.isInteger(attempt) && attempt > 0 && root) pi.registerTool(createMeshControlTool(root, runId, nodeId, attempt));
}