import { createReadStream, readFileSync, writeSync } from "node:fs";

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

import {
  parseWorkerLaunchContractV2,
  validateWorkerChildContext,
  type WorkerChildContext,
} from "./worker-launch.js";
import { classifyProviderFallbackFailure } from "./model-fallback.js";
import type { WorkerLaunchContractV2 } from "./workflow-types.js";

type ProtocolWriter = (value: unknown) => void;

export default function workerChildExtension(pi: ExtensionAPI): void {
  const loaded = loadLaunchContract();
  if (!loaded.ok) {
    pi.on("session_start", (_event, ctx) => refuseWithoutProtocol(loaded.error, ctx));
    pi.on("before_agent_start", (_event, ctx) => refuseWithoutProtocol(loaded.error, ctx));
    return;
  }
  const contract = loaded.contract;
  let writeProtocol: ProtocolWriter;
  try {
    writeProtocol = createProtocolWriter(contract);
  } catch (error) {
    pi.on("session_start", (_event, ctx) => refuseWithoutProtocol(error, ctx));
    pi.on("before_agent_start", (_event, ctx) => refuseWithoutProtocol(error, ctx));
    return;
  }
  let ready = false;
  // The inherited liveness descriptor is single-use for this process: releasing it closes fd 3,
  // so a later re-attach could never succeed. Track the state explicitly instead of detaching
  // and reopening, which would surface as a silent shutdown rather than a visible refusal.
  const liveness = createParentLivenessGuard(contract);

  pi.registerTool({
    name: "message_worker",
    label: "Message Worker",
    description: "Request one directionally granted peer message through the parent supervisor.",
    parameters: Type.Object({
      target: Type.String({ minLength: 1, maxLength: 256 }),
      message: Type.String({ minLength: 1, maxLength: 32_768 }),
      correlation: Type.String({ minLength: 1, maxLength: 256 }),
    }),
    async execute(_id, params) {
      writeProtocol({
        type: "pi_amq_worker_peer_message",
        runId: contract.runId,
        workerId: contract.workerId,
        attemptId: contract.attemptId,
        workerAttemptGeneration: contract.workerAttemptGeneration,
        supervisorGeneration: contract.supervisorGeneration,
        target: params.target,
        message: params.message,
        correlation: params.correlation,
      });
      const details = { status: "requested", target: params.target, correlation: params.correlation };
      return { content: [{ type: "text" as const, text: JSON.stringify(details) }], details };
    },
  });

  pi.registerTool({
    name: "wait_for_input",
    label: "Wait for Input",
    description: "Declare a typed correlated wait without inferring waiting from prose or silence.",
    parameters: Type.Object({
      reason: Type.String({ minLength: 1, maxLength: 2_048 }),
      correlation: Type.String({ minLength: 1, maxLength: 256 }),
    }),
    async execute(_id, params, _signal, _onUpdate, ctx) {
      writeProtocol({
        type: "pi_amq_worker_wait",
        runId: contract.runId,
        workerId: contract.workerId,
        attemptId: contract.attemptId,
        workerAttemptGeneration: contract.workerAttemptGeneration,
        supervisorGeneration: contract.supervisorGeneration,
        reason: params.reason,
        correlation: params.correlation,
      });
      // A typed wait is the terminal action for the current RPC turn. If the tool
      // returns normally, Pi asks the model to continue after the tool result and
      // the model can immediately declare a second wait while the run projection
      // is already blocked. Abort only the active turn; the RPC process remains
      // alive and accepts the supervisor's next prompt.
      ctx.abort();
      const details = {
        status: "declared",
        reason: params.reason,
        correlation: params.correlation,
      };
      return { content: [{ type: "text" as const, text: JSON.stringify(details) }], details };
    },
  });

  pi.on("session_start", (_event, ctx) => {
    try {
      validateWorkerChildContext(contract, childContext(ctx));
      liveness.attach(ctx);
      const sessionFile = ctx.sessionManager.getSessionFile();
      if (!sessionFile) throw new Error("Pi worker session file is unavailable");
      writeProtocol({
        type: "pi_amq_worker_ready",
        runId: contract.runId,
        workerId: contract.workerId,
        attemptId: contract.attemptId,
        workerAttemptGeneration: contract.workerAttemptGeneration,
        supervisorGeneration: contract.supervisorGeneration,
        sessionFile,
      });
      ready = true;
    } catch (error) {
      refuse(error, ctx, writeProtocol);
    }
  });

  pi.on("before_agent_start", (event, ctx) => {
    try {
      if (!ready) throw new Error("Worker launch parity was not established");
      validateWorkerChildContext(contract, childContext(ctx, event.systemPromptOptions.selectedTools));
    } catch (error) {
      refuse(error, ctx, writeProtocol);
    }
  });

  pi.on("model_select", (_event, ctx) => revalidateOrRefuse(contract, ctx, writeProtocol));
  pi.on("thinking_level_select", (_event, ctx) => revalidateOrRefuse(contract, ctx, writeProtocol));
  pi.on("session_shutdown", () => liveness.release());
  pi.on("agent_settled", () => {
    if (!ready) return;
    writeProtocol({
      type: "pi_amq_worker_settled",
      runId: contract.runId,
      workerId: contract.workerId,
      attemptId: contract.attemptId,
      workerAttemptGeneration: contract.workerAttemptGeneration,
      supervisorGeneration: contract.supervisorGeneration,
    });
  });
  pi.on("message_end", (event, ctx) => {
    const message = event.message;
    if (message.role !== "assistant") return;
    const reason = classifyProviderFallbackFailure(message.stopReason, message.errorMessage);
    if (!reason) return;
    writeProtocol({
      type: "pi_amq_worker_provider_failure",
      runId: contract.runId,
      workerId: contract.workerId,
      attemptId: contract.attemptId,
      workerAttemptGeneration: contract.workerAttemptGeneration,
      supervisorGeneration: contract.supervisorGeneration,
      reason,
    });
    ctx.shutdown();
  });
}

function loadLaunchContract():
  | { readonly ok: true; readonly contract: WorkerLaunchContractV2 }
  | { readonly ok: false; readonly error: Error } {
  try {
    const path = process.env["PI_AMQ_WORKER_LAUNCH_CONTRACT"];
    if (!path) throw new Error("PI_AMQ_WORKER_LAUNCH_CONTRACT is required");
    const value = JSON.parse(readFileSync(path, "utf8")) as unknown;
    return { ok: true, contract: parseWorkerLaunchContractV2(value) };
  } catch (error) {
    return { ok: false, error: error instanceof Error ? error : new Error(String(error)) };
  }
}

function childContext(ctx: ExtensionContext, tools?: readonly string[]): WorkerChildContext {
  return {
    mode: ctx.mode,
    cwd: ctx.cwd,
    model: ctx.model ? { provider: ctx.model.provider, id: ctx.model.id } : undefined,
    thinkingLevel: ctx.thinkingLevel,
    projectTrusted: ctx.isProjectTrusted(),
    ...(tools === undefined ? {} : { tools }),
    environmentKeys: Object.keys(process.env),
  };
}

function revalidateOrRefuse(
  contract: WorkerLaunchContractV2,
  ctx: ExtensionContext,
  writeProtocol: ProtocolWriter,
): void {
  try {
    validateWorkerChildContext(contract, childContext(ctx));
  } catch (error) {
    refuse(error, ctx, writeProtocol);
  }
}

interface ParentLivenessGuard {
  /** Idempotent: a repeated session_start keeps the existing stream instead of reopening fd 3. */
  attach(ctx: ExtensionContext): void;
  /** Closes the descriptor so a settled one-shot child can exit. The guard stays released. */
  release(): void;
}

function createParentLivenessGuard(contract: WorkerLaunchContractV2): ParentLivenessGuard {
  let state: "unattached" | "attached" | "released" = "unattached";
  let detach = (): void => {};
  return {
    attach(ctx) {
      if (state === "attached") return;
      if (state === "released") {
        throw new Error("Parent-liveness descriptor was already released and cannot be reattached");
      }
      const configured = Number(process.env["PI_AMQ_PARENT_LIVENESS_FD"]);
      if (configured !== contract.parentLivenessFd) {
        throw new Error("Parent-liveness descriptor does not match the launch contract");
      }
      const stream = createReadStream("", { fd: configured, autoClose: true });
      const requestShutdown = (): void => {
        try {
          ctx.shutdown();
        } catch (error) {
          console.error(`pi-amq parent-liveness shutdown failed: ${errorMessage(error)}`);
        }
      };
      stream.on("end", requestShutdown);
      stream.on("error", requestShutdown);
      stream.resume();
      detach = () => {
        stream.off("end", requestShutdown);
        stream.off("error", requestShutdown);
        stream.destroy();
      };
      state = "attached";
    },
    release() {
      if (state !== "attached") return;
      detach();
      detach = (): void => {};
      state = "released";
    },
  };
}

function refuse(error: unknown, ctx: ExtensionContext, writeProtocol: ProtocolWriter): void {
  try {
    writeProtocol({
      type: "pi_amq_worker_refused",
      error: error instanceof Error ? error.message : String(error),
    });
  } catch (protocolError) {
    console.error(`pi-amq worker protocol refusal failed: ${errorMessage(protocolError)}`);
  }
  ctx.shutdown();
}

function refuseWithoutProtocol(error: unknown, ctx: ExtensionContext): void {
  console.error(`pi-amq worker launch refused: ${errorMessage(error)}`);
  ctx.shutdown();
}

function createProtocolWriter(contract: WorkerLaunchContractV2): ProtocolWriter {
  const configured = Number(process.env["PI_AMQ_PROTOCOL_FD"]);
  if (configured !== contract.protocolFd) {
    throw new Error("Worker protocol descriptor does not match the launch contract");
  }
  return (value) => writeProtocolLine(contract.protocolFd, value);
}

function writeProtocolLine(fd: number, value: unknown): void {
  const line = Buffer.from(`${JSON.stringify(value)}\n`, "utf8");
  let offset = 0;
  while (offset < line.byteLength) {
    const written = writeSync(fd, line, offset, line.byteLength - offset);
    if (written <= 0) throw new Error("Worker protocol descriptor accepted no bytes");
    offset += written;
  }
}

function errorMessage(error: unknown): string {
  return error instanceof Error ? error.message : String(error);
}