import { randomUUID } from "node:crypto";
import { existsSync, readFileSync } from "node:fs";
import { join } from "node:path";

import { StringEnum } from "@earendil-works/pi-ai";
import type {
  ExtensionAPI,
  ExtensionContext
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

import {
  appendCheckpoint,
  appendDelegate,
  appendRoute,
  createRunState,
  isTerminalStatus,
  normalizeRunState,
  transitionRunState,
  updateDelegateStatus
} from "../lib/run-state.mjs";
import { resolveRoute, resolveRoutingConfig } from "../lib/routing.mjs";
import {
  assertCheckpointAllowed,
  assertDelegateAllowed
} from "../lib/workflow-guards.mjs";
import { DELEGATE_ACCEPTANCE } from "../lib/delegation-policy.mjs";

const STATE_TYPE = "adaptive-orchestrator-state";
const MAX_REQUEST_CHARS = 8_000;
const MAX_SUMMARY_CHARS = 4_000;
const RPC_REQUEST = "subagents:rpc:v1:request";
const RPC_REPLY_PREFIX = "subagents:rpc:v1:reply:";
const ASYNC_COMPLETE = "subagent:async-complete";

const RouteParams = Type.Object({
  role: StringEnum(["planner", "final-review"] as const),
  complexity: StringEnum(["easy", "medium", "complex"] as const),
  rationale: Type.String({ minLength: 1, maxLength: 2_000 })
});

const DelegateParams = Type.Object({
  role: StringEnum(["worker", "reviewer"] as const),
  purpose: StringEnum([
    "implementation",
    "fix",
    "initial-review",
    "re-review"
  ] as const),
  complexity: StringEnum(["easy", "medium", "complex"] as const),
  rationale: Type.String({ minLength: 1, maxLength: 2_000 }),
  workId: Type.String({ minLength: 1, maxLength: 120 }),
  retry: Type.Optional(Type.Boolean()),
  task: Type.String({ minLength: 1, maxLength: MAX_REQUEST_CHARS })
});

const CheckpointParams = Type.Object({
  phase: StringEnum([
    "plan-ready",
    "workers-complete",
    "review-complete"
  ] as const),
  summary: Type.String({ minLength: 1, maxLength: MAX_SUMMARY_CHARS })
});

const HandoffParams = Type.Object({
  summary: Type.String({ minLength: 1, maxLength: MAX_SUMMARY_CHARS })
});

const escapePromptText = (value: string) =>
  value
    .replaceAll("&", "&")
    .replaceAll("<", "&lt;")
    .replaceAll(">", "&gt;");

const isRecord = (value: unknown): value is Record<string, unknown> =>
  Boolean(value) && typeof value === "object" && !Array.isArray(value);

const configPath = (ctx: ExtensionContext) =>
  join(ctx.cwd, ".pi", "adaptive-orchestrator.json");

const loadConfig = (ctx: ExtensionContext) => {
  if (!ctx.isProjectTrusted() || !existsSync(configPath(ctx))) {
    return resolveRoutingConfig();
  }
  return resolveRoutingConfig(
    JSON.parse(readFileSync(configPath(ctx), "utf8"))
  );
};

const statusSummary = (state: ReturnType<typeof normalizeRunState>) => {
  if (!state) return "No adaptive orchestration run is active.";
  const route = state.routes.at(-1);
  const running = state.delegates.filter(
    (delegate) => delegate.status === "running"
  );
  return [
    `Adaptive orchestrator: ${state.status}`,
    `Request: ${state.request}`,
    route
      ? `Latest route: ${route.role} → ${route.model}:${route.thinking}`
      : "Latest route: pending",
    `Delegates: ${state.delegates.length} total, ${running.length} running`
  ].join("\n");
};

const planningInstructions = (
  request: string
) => `You are the planning coordinator for a human-approved adaptive orchestration run.

<untrusted_request>
${escapePromptText(request)}
</untrusted_request>

Start at Sol medium effort. First assess complexity from scope, uncertainty, dependency breadth, safety risk, and validation burden. Call orchestration_route with role="planner" and your evidence-backed complexity decision; it will move Sol to the approved planning effort (Pi calls “extra high” xhigh). Then produce a detailed plan: goals/non-goals, assumptions, affected areas, ordered work slices, exact validation, risks, and a future worker/reviewer complexity recommendation for every slice.

Do not implement, edit files, run shell commands, or launch subagents in this phase. End by calling orchestration_checkpoint with phase="plan-ready" and a concise plan summary. The extension will pause for explicit human approval.`;

const executionInstructions = (
  request: string
) => `Execute the approved plan as the coordinator only.

<untrusted_request>
${escapePromptText(request)}
</untrusted_request>

Never directly edit files, run shell commands, or call subagent. Use orchestration_delegate for every child: it routes the exact model, rejects pre-approval work, tracks completion, prevents concurrent active-worktree writers, and prevents child nesting. After every delegate, call wait with its returned run id before the next checkpoint or delegate. Workers are serialized. Reviewers are read-only and begin only after orchestration_checkpoint phase="workers-complete".

Delegate implementation slices, wait for their completion in /orchestrate status, then record workers-complete. Delegate one initial-review. If it finds actionable blockers, delegate exactly one fix worker and then one re-review. Record review-complete only after all required delegate results are complete. Route final-review to move yourself to the approved Sol effort, independently inspect the resulting evidence, then call orchestration_handoff. Do not deploy, merge, or claim human approval.`;

export default function adaptiveOrchestratorExtension(pi: ExtensionAPI) {
  let state: ReturnType<typeof normalizeRunState> = null;
  let routingConfig = resolveRoutingConfig();
  let configError: string | undefined;
  let runtimeContext: ExtensionContext | null = null;
  const completionEvents = new Map<string, unknown>();

  const persist = () => pi.appendEntry(STATE_TYPE, { version: 1, state });

  const updateUi = (ctx = runtimeContext) => {
    if (!ctx?.hasUI) return;
    ctx.ui.setStatus(
      "adaptive-orchestrator",
      state
        ? ctx.ui.theme.fg(
            isTerminalStatus(state.status) ? "success" : "accent",
            `Orchestrator: ${state.status}`
          )
        : undefined
    );
  };

  const show = (
    ctx: ExtensionContext,
    message: string,
    level: "info" | "warning" | "error" = "info"
  ) => {
    if (ctx.hasUI) ctx.ui.notify(message.split("\n")[0] ?? message, level);
    pi.sendMessage(
      { customType: "adaptive-orchestrator", content: message, display: true },
      { triggerTurn: false }
    );
  };

  const reloadState = (ctx: ExtensionContext) => {
    runtimeContext = ctx;
    state = ctx.sessionManager
      .getBranch()
      .reduce<ReturnType<typeof normalizeRunState>>((current, entry) => {
        if (entry.type !== "custom" || entry.customType !== STATE_TYPE) {
          return current;
        }
        const data = entry.data as { state?: unknown } | undefined;
        return normalizeRunState(data?.state);
      }, null);
    try {
      routingConfig = loadConfig(ctx);
      configError = undefined;
    } catch (error) {
      routingConfig = resolveRoutingConfig();
      configError = error instanceof Error ? error.message : String(error);
    }
    updateUi(ctx);
  };

  const requireConfig = () => {
    if (configError) {
      throw new Error(`Invalid .pi/adaptive-orchestrator.json: ${configError}`);
    }
    return routingConfig;
  };

  const setParentRoute = async (
    ctx: ExtensionContext,
    route: ReturnType<typeof resolveRoute>
  ) => {
    const [provider, id] = route.model.split("/", 2);
    const model = ctx.modelRegistry.find(provider, id);
    if (!model) {
      throw new Error(
        `Configured model is unavailable in this Pi session: ${route.model}`
      );
    }
    if (!(await pi.setModel(model))) {
      throw new Error(
        `No usable credentials for configured model: ${route.model}`
      );
    }
    pi.setThinkingLevel(route.thinking);
  };

  const callSubagentRpc = async (
    params: Record<string, unknown>
  ): Promise<Record<string, unknown>> => {
    const requestId = randomUUID();
    return new Promise((resolve, reject) => {
      const replyEvent = `${RPC_REPLY_PREFIX}${requestId}`;
      const timeout = setTimeout(() => {
        cleanup();
        reject(
          new Error("pi-subagents did not acknowledge the delegation request.")
        );
      }, 10_000);
      const unsubscribe = pi.events.on(replyEvent, (raw) => {
        cleanup();
        if (!isRecord(raw) || raw.success !== true) {
          const message =
            isRecord(raw) &&
            isRecord(raw.error) &&
            typeof raw.error.message === "string"
              ? raw.error.message
              : "pi-subagents rejected the delegation request.";
          reject(new Error(message));
          return;
        }
        resolve(isRecord(raw.data) ? raw.data : {});
      });
      const cleanup = () => {
        clearTimeout(timeout);
        if (typeof unsubscribe === "function") unsubscribe();
      };
      pi.events.emit(RPC_REQUEST, {
        version: 1,
        requestId,
        method: "spawn",
        params,
        source: { extension: "pi-adaptive-orchestrator" }
      });
    });
  };

  const delegateStatus = (raw: unknown) => {
    const resultStatuses =
      isRecord(raw) && Array.isArray(raw.results)
        ? raw.results
            .filter(isRecord)
            .map((result) => result.status)
            .filter((status): status is string => typeof status === "string")
        : [];
    if (resultStatuses.includes("failed")) return "failed";
    if (resultStatuses.includes("paused")) return "paused";
    return "completed";
  };

  const coordinatorTools = new Set([
    "read",
    "grep",
    "find",
    "ls",
    "wait",
    "orchestration_route",
    "orchestration_delegate",
    "orchestration_checkpoint",
    "orchestration_handoff"
  ]);

  pi.on("session_start", async (_event, ctx) => reloadState(ctx));
  pi.on("session_tree", async (_event, ctx) => reloadState(ctx));
  pi.on("session_shutdown", async (_event, ctx) => {
    updateUi(ctx);
    runtimeContext = null;
  });

  pi.events.on(ASYNC_COMPLETE, (raw) => {
    if (!isRecord(raw) || typeof raw.runId !== "string") return;
    const delegate = state?.delegates.find(
      (entry) => entry.runId === raw.runId
    );
    if (!state || !delegate) {
      completionEvents.set(raw.runId, raw);
      return;
    }
    if (delegate.status !== "running") return;
    state = updateDelegateStatus(state, raw.runId, delegateStatus(raw));
    persist();
    updateUi();
  });

  pi.on("before_agent_start", async (event) => {
    if (!state || isTerminalStatus(state.status)) return;
    const instruction =
      state.status === "planning"
        ? planningInstructions(state.request)
        : state.status === "awaiting-approval"
          ? "The adaptive plan is awaiting human approval. Do not implement or launch subagents; answer questions about the plan only."
          : executionInstructions(state.request);
    return { systemPrompt: `${event.systemPrompt}\n\n${instruction}` };
  });

  pi.on("tool_call", async (event) => {
    if (!state || isTerminalStatus(state.status)) return;
    if (!coordinatorTools.has(event.toolName)) {
      return {
        block: true,
        reason:
          "The adaptive coordinator is restricted to read-only inspection and orchestration tools. Delegate all implementation through orchestration_delegate."
      };
    }
  });

  pi.registerCommand("orchestrate", {
    description:
      "Plan a task with effort routing, approve it, monitor delegated work, or cancel it",
    getArgumentCompletions: (prefix) =>
      ["approve", "status", "cancel"]
        .filter((value) => value.startsWith(prefix.trim()))
        .map((value) => ({ value, label: value })),
    handler: async (args, ctx) => {
      const input = args.trim();
      try {
        if (!input || input === "status") {
          show(ctx, statusSummary(state));
          return;
        }
        if (input === "cancel") {
          if (!state || isTerminalStatus(state.status)) {
            throw new Error("No active orchestration run can be cancelled.");
          }
          state = transitionRunState(state, "cancelled");
          persist();
          updateUi(ctx);
          show(ctx, statusSummary(state), "warning");
          return;
        }
        if (input === "approve") {
          if (!state || state.status !== "awaiting-approval") {
            throw new Error(
              "Approval is available only after the planner checkpoint is ready."
            );
          }
          state = transitionRunState(state, "executing");
          persist();
          updateUi(ctx);
          pi.sendUserMessage(executionInstructions(state.request));
          await ctx.waitForIdle();
          return;
        }
        if (input.length > MAX_REQUEST_CHARS) {
          throw new Error(
            `Request exceeds ${MAX_REQUEST_CHARS.toLocaleString()} characters.`
          );
        }
        if (state && !isTerminalStatus(state.status)) {
          throw new Error(
            "An orchestration run is already active. Use /orchestrate status, approve, or cancel."
          );
        }
        const initialRoute = resolveRoute({
          role: "planner",
          complexity: "easy",
          config: requireConfig()
        });
        await setParentRoute(ctx, initialRoute);
        state = createRunState({ id: randomUUID(), request: input });
        persist();
        updateUi(ctx);
        pi.sendUserMessage(planningInstructions(state.request));
        await ctx.waitForIdle();
      } catch (error) {
        show(
          ctx,
          error instanceof Error ? error.message : String(error),
          "error"
        );
      }
    }
  });

  pi.registerTool({
    name: "orchestration_route",
    label: "Orchestration Route",
    description:
      "Record the evidence-backed Sol planning or final-review complexity decision and change the parent effort.",
    promptSnippet:
      "Route Sol planning and final review to the validated effort",
    promptGuidelines: [
      "Use orchestration_route only for the parent planner and final review; orchestration_delegate routes every child itself."
    ],
    parameters: RouteParams,
    async execute(_id, params, _signal, _onUpdate, ctx) {
      if (!state || isTerminalStatus(state.status)) {
        throw new Error("No active orchestration run exists.");
      }
      const allowed =
        (params.role === "planner" && state.status === "planning") ||
        (params.role === "final-review" &&
          state.status === "reviewing" &&
          state.checkpoints.some(
            (checkpoint) => checkpoint.phase === "review-complete"
          ));
      if (!allowed) {
        throw new Error(
          `Role '${params.role}' is not allowed while orchestration is ${state.status}.`
        );
      }
      const route = resolveRoute({
        role: params.role,
        complexity: params.complexity,
        config: requireConfig()
      });
      state = appendRoute(state, route, params.rationale);
      await setParentRoute(ctx, route);
      persist();
      updateUi(ctx);
      return {
        content: [{ type: "text", text: JSON.stringify(route, null, 2) }],
        details: route
      };
    }
  });

  pi.registerTool({
    name: "orchestration_delegate",
    label: "Orchestration Delegate",
    description:
      "Route and start one tracked worker or reviewer through pi-subagents with enforced writer serialization and child constraints.",
    promptSnippet:
      "Delegate a tracked, effort-routed worker or reviewer without direct subagent calls",
    promptGuidelines: [
      "Use orchestration_delegate for every worker and reviewer; do not call subagent directly during an adaptive orchestration run."
    ],
    parameters: DelegateParams,
    async execute(_id, params, _signal, _onUpdate, ctx) {
      if (!state || isTerminalStatus(state.status)) {
        throw new Error("No active orchestration run exists.");
      }
      assertDelegateAllowed(
        state,
        params.role,
        params.purpose,
        params.workId,
        params.retry === true
      );

      const route = resolveRoute({
        role: params.role,
        complexity: params.complexity,
        config: requireConfig()
      });
      state = appendRoute(state, route, params.rationale);
      persist();
      const childContract =
        params.role === "worker"
          ? "You are the sole active-worktree writer. Make only this approved change, run focused validation, and do not spawn subagents."
          : "You are a read-only independent reviewer. Do not modify project/source files or spawn subagents; report evidence-backed findings with severity and file/line references.";
      const response = await callSubagentRpc({
        agent: route.agent,
        task: `${childContract}\n\nApproved task:\n${params.task}`,
        cwd: ctx.cwd,
        context: "fresh",
        model: route.subagentModel,
        async: true,
        acceptance: DELEGATE_ACCEPTANCE
      });
      const details = isRecord(response.details) ? response.details : {};
      const runId =
        typeof details.runId === "string" ? details.runId : undefined;
      if (!runId) {
        throw new Error(
          "pi-subagents started a child without a trackable run id."
        );
      }
      state = appendDelegate(state, {
        runId,
        workId: params.workId,
        retry: params.retry === true,
        role: params.role,
        purpose: params.purpose,
        model: route.subagentModel,
        task: params.task
      });
      const completedBeforeTracking = completionEvents.get(runId);
      if (completedBeforeTracking) {
        completionEvents.delete(runId);
        state = updateDelegateStatus(
          state,
          runId,
          delegateStatus(completedBeforeTracking)
        );
      }
      persist();
      updateUi(ctx);
      return {
        content: [
          { type: "text", text: `Started ${params.purpose} delegate ${runId}.` }
        ],
        details: { runId, route }
      };
    }
  });

  pi.registerTool({
    name: "orchestration_checkpoint",
    label: "Orchestration Checkpoint",
    description:
      "Persist the plan, completed workers, or completed review checkpoint for the active adaptive run.",
    promptSnippet:
      "Record plan, worker-completion, and review-completion checkpoints",
    promptGuidelines: [
      "Use orchestration_checkpoint with plan-ready before asking the human to approve work; workers and reviewers must be completed in the tracked state before their checkpoints."
    ],
    parameters: CheckpointParams,
    async execute(_id, params, _signal, _onUpdate, ctx) {
      if (!state || isTerminalStatus(state.status)) {
        throw new Error("No active orchestration run exists.");
      }
      const nextStatus = assertCheckpointAllowed(state, params.phase);
      state = appendCheckpoint(state, params.phase, params.summary);
      if (state.status !== nextStatus)
        state = transitionRunState(state, nextStatus);
      persist();
      updateUi(ctx);
      const message =
        params.phase === "plan-ready"
          ? "Plan is persisted and awaiting explicit human approval. Ask the human to run /orchestrate approve."
          : `Checkpoint recorded: ${params.phase}.`;
      return {
        content: [{ type: "text", text: message }],
        details: { state, summary: params.summary }
      };
    }
  });

  pi.registerTool({
    name: "orchestration_handoff",
    label: "Orchestration Handoff",
    description:
      "Finish adaptive orchestration after the tracked review loop and routed final Sol review, then hand evidence to a human.",
    promptSnippet:
      "Mark adaptive orchestration ready for human review after final Sol review",
    promptGuidelines: [
      "Use orchestration_handoff only after review-complete and an evidence-backed final-review route; it never deploys, merges, or substitutes for human approval."
    ],
    parameters: HandoffParams,
    async execute(_id, params, _signal, _onUpdate, ctx) {
      if (!state || state.status !== "reviewing") {
        throw new Error("Human handoff requires an active reviewing state.");
      }
      if (
        !state.checkpoints.some(
          (checkpoint) => checkpoint.phase === "review-complete"
        )
      ) {
        throw new Error(
          "Complete and record the independent review before human handoff."
        );
      }
      if (!state.routes.some((route) => route.role === "final-review")) {
        throw new Error("Route and perform final-review before human handoff.");
      }
      state = appendCheckpoint(state, "human-handoff", params.summary);
      state = transitionRunState(state, "ready-for-human");
      persist();
      updateUi(ctx);
      const message = `Adaptive orchestration is ready for human review.\n${params.summary}`;
      show(ctx, message);
      return { content: [{ type: "text", text: message }], details: { state } };
    }
  });
}