import { isAbsolute, normalize, relative, sep } from "node:path";

import type {
  AttemptId,
  ModelSpec,
  WorkerContract,
  WorkerLaunchContract,
  WorkerLaunchContractV1,
  WorkerLaunchContractV2,
} from "./workflow-types.js";
import { parseBudgetLimit } from "./budget.js";

const FORBIDDEN_EXACT_ENVIRONMENT_KEYS = new Set([
  "PI_SESSION_ID",
  "PI_SESSION_FILE",
  "PI_PROVIDER",
  "PI_MODEL",
  "PI_REASONING_LEVEL",
]);

export interface WorkerLaunchAuthority {
  readonly piExecutable: string;
  readonly childExtensionPath: string;
  readonly grantedCapabilities: readonly string[];
  readonly capabilityTools: Readonly<Record<string, readonly string[]>>;
  readonly availableModels: readonly {
    readonly provider: string;
    readonly id: string;
    readonly thinkingLevels: readonly string[];
    readonly defaultThinkingLevel?: string;
  }[];
  readonly projectTrust: "trusted" | "untrusted" | "unresolved";
  readonly requiresProjectResources: boolean;
  readonly inheritedEnvironment: Readonly<Record<string, string | undefined>>;
}

export interface WorkerLaunchAttemptIdentity {
  readonly attemptId: AttemptId;
  readonly workerAttemptGeneration: string;
  readonly supervisorGeneration: string;
}

export interface PiProcessLaunchCommand {
  readonly executable: string;
  readonly args: readonly string[];
  readonly cwd: string;
  readonly env: Readonly<Record<string, string>>;
  readonly parentLivenessFd: 3;
  readonly protocolFd: 4;
  readonly launchContract: WorkerLaunchContractV2;
  readonly launchContractPath: string;
}

export interface ResolvedWorkerLaunch {
  readonly contract: WorkerLaunchContractV2;
  readonly command: PiProcessLaunchCommand;
}

export interface WorkerChildContext {
  readonly mode: "tui" | "rpc" | "json" | "print";
  readonly cwd: string;
  readonly model: { readonly provider: string; readonly id: string } | undefined;
  readonly thinkingLevel: string | undefined;
  readonly projectTrusted: boolean;
  readonly tools?: readonly string[];
  readonly environmentKeys: readonly string[];
}

export class WorkerLaunchValidationError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "WorkerLaunchValidationError";
  }
}

export function parseWorkerLaunchContractV1(value: unknown): WorkerLaunchContractV1 {
  const parsed = parseWorkerLaunchContractVersion(value, 1);
  if (parsed.schema !== 1) throw new WorkerLaunchValidationError("Worker launch schema must be 1");
  return parsed;
}

export function parseWorkerLaunchContractV2(value: unknown): WorkerLaunchContractV2 {
  const parsed = parseWorkerLaunchContractVersion(value, 2);
  if (parsed.schema !== 2) throw new WorkerLaunchValidationError("Worker launch schema must be 2");
  return parsed;
}

export function parseWorkerLaunchContract(value: unknown): WorkerLaunchContract {
  if (typeof value !== "object" || value === null || Array.isArray(value)) {
    throw new WorkerLaunchValidationError("worker launch contract must be an object");
  }
  const schema = (value as Record<string, unknown>)["schema"];
  if (schema === 1) return parseWorkerLaunchContractV1(value);
  if (schema === 2) return parseWorkerLaunchContractV2(value);
  throw new WorkerLaunchValidationError("Worker launch schema must be 1 or 2");
}

function parseWorkerLaunchContractVersion(
  value: unknown,
  schema: 1 | 2,
): WorkerLaunchContract {
  const object = strictObject(value, [
    "schema", "runId", "workerId", "attemptId", "workerAttemptGeneration",
    "supervisorGeneration", "transportMode", "channel", "depth", "model", "capabilities",
    "tools", "workspace", "sessionDirectory", "projectTrust", "parentLivenessFd", "budget",
    "sideEffectPolicy", "continuedSessionFile", "forkedSessionFile",
    ...(schema === 2 ? ["protocolFd"] : []),
  ], "worker launch contract");
  if (object["schema"] !== schema) {
    throw new WorkerLaunchValidationError(`Worker launch schema must be ${schema}`);
  }
  if (object["transportMode"] !== "send-only") {
    throw new WorkerLaunchValidationError("Worker launch transport mode must be send-only");
  }
  if (object["parentLivenessFd"] !== 3) {
    throw new WorkerLaunchValidationError("Worker parent-liveness descriptor must be 3");
  }
  if (schema === 2 && object["protocolFd"] !== 4) {
    throw new WorkerLaunchValidationError("Worker protocol descriptor must be 4");
  }
  const model = strictObject(object["model"], ["provider", "id", "thinkingLevel"], "worker model");
  const workspace = strictObject(object["workspace"], ["path", "mode", "worktreePath"], "worker workspace");
  const budget = strictObject(
    object["budget"],
    ["maxAttempts", "maxWallClockMs", "maxTokens", "maxCostUsd", "maxOutputBytes"],
    "worker budget",
  );
  const thinkingLevel = requiredString(model["thinkingLevel"], "worker thinking level");
  const worktreePath = workspace["worktreePath"] === undefined
    ? undefined
    : absolutePath(String(workspace["worktreePath"]), "worker worktree path");
  const continuedSessionFile = object["continuedSessionFile"] === undefined
    ? undefined
    : absolutePath(String(object["continuedSessionFile"]), "continued worker session file");
  const forkedSessionFile = object["forkedSessionFile"] === undefined
    ? undefined
    : absolutePath(String(object["forkedSessionFile"]), "forked worker session file");
  const sessionDirectory = absolutePath(
    String(object["sessionDirectory"]),
    "worker session directory",
  );
  if (continuedSessionFile !== undefined && !pathIsInside(sessionDirectory, continuedSessionFile)) {
    throw new WorkerLaunchValidationError(
      "Continued worker session file must be inside the leased session directory",
    );
  }
  if (forkedSessionFile !== undefined && !pathIsInside(sessionDirectory, forkedSessionFile)) {
    throw new WorkerLaunchValidationError(
      "Forked worker session file must be inside the leased session directory",
    );
  }
  if (continuedSessionFile !== undefined && forkedSessionFile !== undefined) {
    throw new WorkerLaunchValidationError("Worker launch cannot continue and fork a session together");
  }
  const fields: Omit<WorkerLaunchContractV1, "schema"> = {
    runId: prefixedString(object["runId"], "run", "run ID") as WorkerLaunchContract["runId"],
    workerId: prefixedString(object["workerId"], "worker", "worker ID") as WorkerLaunchContract["workerId"],
    attemptId: prefixedString(object["attemptId"], "attempt", "attempt ID") as WorkerLaunchContract["attemptId"],
    workerAttemptGeneration: requiredString(
      object["workerAttemptGeneration"],
      "worker attempt generation",
    ),
    supervisorGeneration: requiredString(object["supervisorGeneration"], "supervisor generation"),
    transportMode: "send-only",
    channel: exactString(object["channel"], ["json-result", "rpc"] as const, "worker channel"),
    depth: positiveIntegerValue(object["depth"], "worker depth"),
    model: {
      provider: requiredString(model["provider"], "worker model provider"),
      id: requiredString(model["id"], "worker model ID"),
      thinkingLevel,
    },
    capabilities: stringArray(object["capabilities"], "worker capabilities"),
    tools: stringArray(object["tools"], "worker tools"),
    workspace: {
      path: absolutePath(String(workspace["path"]), "worker workspace path"),
      mode: exactString(workspace["mode"], ["read-only", "mutating"] as const, "workspace mode"),
      ...(worktreePath ? { worktreePath } : {}),
    },
    sessionDirectory,
    ...(continuedSessionFile ? { continuedSessionFile } : {}),
    ...(forkedSessionFile ? { forkedSessionFile } : {}),
    projectTrust: exactString(
      object["projectTrust"],
      ["trusted", "untrusted"] as const,
      "worker project trust",
    ),
    parentLivenessFd: 3,
    budget: {
      maxAttempts: parseBudgetLimit(budget["maxAttempts"], "worker maxAttempts", { integer: true, allowLegacyNumber: true }),
      maxWallClockMs: parseBudgetLimit(budget["maxWallClockMs"], "worker maxWallClockMs", { integer: true, allowLegacyNumber: true }),
      maxTokens: parseBudgetLimit(budget["maxTokens"], "worker maxTokens", { integer: true, allowLegacyNumber: true }),
      maxCostUsd: budget["maxCostUsd"] === undefined
        ? { kind: "unlimited" }
        : parseBudgetLimit(budget["maxCostUsd"], "worker maxCostUsd", { integer: false, allowLegacyNumber: true }),
      maxOutputBytes: budget["maxOutputBytes"] === undefined
        ? { kind: "unlimited" }
        : parseBudgetLimit(budget["maxOutputBytes"], "worker maxOutputBytes", { integer: true, allowLegacyNumber: true }),
    },
    sideEffectPolicy: exactString(
      object["sideEffectPolicy"],
      ["none", "idempotent", "reconcilable", "external"] as const,
      "worker side-effect policy",
    ),
  };
  return schema === 1
    ? { schema: 1, ...fields }
    : { schema: 2, ...fields, protocolFd: 4 };
}

export function resolveWorkerLaunch(input: {
  readonly contract: WorkerContract;
  readonly attempt: WorkerLaunchAttemptIdentity;
  readonly depth: number;
  readonly prompt: string;
  readonly sessionDirectory: string;
  readonly launchContractPath: string;
  readonly continuedSessionFile?: string;
  readonly forkedSessionFile?: string;
  readonly modelOverride?: ModelSpec;
  readonly authority: WorkerLaunchAuthority;
}): ResolvedWorkerLaunch {
  const { contract, attempt, authority } = input;
  const prompt = requireNonEmpty(input.prompt, "worker prompt");
  const depth = positiveInteger(input.depth, "worker depth");
  const sessionDirectory = absolutePath(input.sessionDirectory, "worker session directory");
  const launchContractPath = absolutePath(input.launchContractPath, "worker launch-contract path");
  const continuedSessionFile = input.continuedSessionFile === undefined
    ? undefined
    : absolutePath(input.continuedSessionFile, "continued worker session file");
  const forkedSessionFile = input.forkedSessionFile === undefined
    ? undefined
    : absolutePath(input.forkedSessionFile, "forked worker session file");
  if (continuedSessionFile !== undefined
    && !pathIsInside(sessionDirectory, continuedSessionFile)) {
    throw new WorkerLaunchValidationError(
      "Continued worker session file must be inside the leased session directory",
    );
  }
  if (forkedSessionFile !== undefined && !pathIsInside(sessionDirectory, forkedSessionFile)) {
    throw new WorkerLaunchValidationError(
      "Forked worker session file must be inside the leased session directory",
    );
  }
  if (continuedSessionFile !== undefined && forkedSessionFile !== undefined) {
    throw new WorkerLaunchValidationError("Worker launch cannot continue and fork a session together");
  }
  const executable = absolutePath(authority.piExecutable, "pi executable");
  const childExtensionPath = absolutePath(authority.childExtensionPath, "worker child extension");
  const granted = new Set(authority.grantedCapabilities);
  for (const capability of contract.capabilities) {
    const authorityCapability = capability.startsWith("message_worker:")
      ? "message_worker"
      : capability;
    if (!granted.has(authorityCapability)) {
      throw new WorkerLaunchValidationError(`Worker capability is not granted: ${capability}`);
    }
    if (!(authorityCapability in authority.capabilityTools)) {
      throw new WorkerLaunchValidationError(`Worker capability has no tool projection: ${capability}`);
    }
  }
  if (contract.channel === "json-result"
    && contract.capabilities.some((capability) =>
      capability === "worker.steer"
      || capability === "driver-inbox.forward"
      || capability.startsWith("message_worker:"))) {
    throw new WorkerLaunchValidationError(
      "Live steering and inbound forwarding require the RPC channel",
    );
  }
  const requestedModel = input.modelOverride ?? contract.model;
  const selectedModel = authority.availableModels.find(
    (model) => model.provider === requestedModel.provider && model.id === requestedModel.id,
  );
  if (!selectedModel) {
    throw new WorkerLaunchValidationError(
      `Worker model is not available: ${requestedModel.provider}/${requestedModel.id}`,
    );
  }
  const thinkingLevel = requestedModel.thinkingLevel ?? selectedModel.defaultThinkingLevel;
  if (thinkingLevel === undefined || !selectedModel.thinkingLevels.includes(thinkingLevel)) {
    throw new WorkerLaunchValidationError(
      `Worker thinking level is unavailable or unresolved for ${requestedModel.provider}/${requestedModel.id}`,
    );
  }
  if (authority.requiresProjectResources && authority.projectTrust !== "trusted") {
    throw new WorkerLaunchValidationError(
      "Required project resources cannot be loaded because project trust is unresolved or denied",
    );
  }
  const projectTrust = authority.requiresProjectResources ? "trusted" : "untrusted";
  const tools = [...new Set([
    ...contract.capabilities.flatMap(
      (capability) => authority.capabilityTools[
        capability.startsWith("message_worker:") ? "message_worker" : capability
      ] ?? [],
    ),
    ...(contract.channel === "rpc" ? ["wait_for_input"] : []),
  ])].sort();
  const launchContract: WorkerLaunchContractV2 = {
    schema: 2,
    runId: contract.runId,
    workerId: contract.workerId,
    attemptId: attempt.attemptId,
    workerAttemptGeneration: requireNonEmpty(
      attempt.workerAttemptGeneration,
      "worker attempt generation",
    ),
    supervisorGeneration: requireNonEmpty(attempt.supervisorGeneration, "supervisor generation"),
    transportMode: "send-only",
    channel: contract.channel,
    depth,
    model: {
      provider: requestedModel.provider,
      id: requestedModel.id,
      thinkingLevel,
    },
    capabilities: [...contract.capabilities],
    tools,
    workspace: contract.workspace,
    sessionDirectory,
    ...(continuedSessionFile ? { continuedSessionFile } : {}),
    ...(forkedSessionFile ? { forkedSessionFile } : {}),
    projectTrust,
    parentLivenessFd: 3,
    protocolFd: 4,
    budget: contract.budget,
    sideEffectPolicy: contract.sideEffectPolicy,
  };
  const args = buildPiArguments(
    launchContract,
    childExtensionPath,
    authority.requiresProjectResources,
    prompt,
  );
  return {
    contract: launchContract,
    command: {
      executable,
      args,
      cwd: absolutePath(contract.workspace.worktreePath ?? contract.workspace.path, "worker cwd"),
      env: childEnvironment(authority.inheritedEnvironment, launchContractPath),
      parentLivenessFd: 3,
      protocolFd: 4,
      launchContract,
      launchContractPath,
    },
  };
}

export function validateWorkerChildContext(
  contract: WorkerLaunchContract,
  context: WorkerChildContext,
): void {
  const expectedMode = contract.channel === "json-result" ? "json" : "rpc";
  if (context.mode !== expectedMode) {
    throw new WorkerLaunchValidationError(
      `Worker mode parity mismatch: expected ${expectedMode}, observed ${context.mode}`,
    );
  }
  if (normalize(context.cwd) !== normalize(contract.workspace.worktreePath ?? contract.workspace.path)) {
    throw new WorkerLaunchValidationError("Worker cwd parity mismatch");
  }
  if (!context.model
    || context.model.provider !== contract.model.provider
    || context.model.id !== contract.model.id) {
    throw new WorkerLaunchValidationError("Worker model parity mismatch");
  }
  if (context.thinkingLevel !== contract.model.thinkingLevel) {
    throw new WorkerLaunchValidationError("Worker thinking-level parity mismatch");
  }
  if (context.projectTrusted !== (contract.projectTrust === "trusted")) {
    throw new WorkerLaunchValidationError("Worker project-trust parity mismatch");
  }
  const forbidden = context.environmentKeys.filter(isReceiverAuthorityEnvironmentKey);
  if (forbidden.length > 0) {
    throw new WorkerLaunchValidationError(
      `Worker inherited receiver authority: ${forbidden.sort().join(", ")}`,
    );
  }
  if (context.tools !== undefined
    && canonicalStrings(context.tools) !== canonicalStrings(contract.tools)) {
    throw new WorkerLaunchValidationError("Worker tool parity mismatch");
  }
}

export function isReceiverAuthorityEnvironmentKey(key: string): boolean {
  return key.startsWith("AM_")
    || key.startsWith("ORCH_")
    || FORBIDDEN_EXACT_ENVIRONMENT_KEYS.has(key);
}

function buildPiArguments(
  contract: WorkerLaunchContract,
  childExtensionPath: string,
  requiresProjectResources: boolean,
  prompt: string,
): readonly string[] {
  const args = [
    "--provider", contract.model.provider,
    "--model", contract.model.id,
    "--thinking", contract.model.thinkingLevel!,
    "--session-dir", contract.sessionDirectory,
    ...((contract.continuedSessionFile ?? contract.forkedSessionFile)
      ? ["--session", (contract.continuedSessionFile ?? contract.forkedSessionFile)!]
      : []),
    "--name", `pi-amq:${contract.workerId}`,
    requiresProjectResources ? "--approve" : "--no-approve",
    "--no-extensions",
    "-e", childExtensionPath,
  ];
  if (!requiresProjectResources) {
    args.push("--no-skills", "--no-prompt-templates", "--no-context-files");
  }
  if (contract.tools.length === 0) args.push("--no-tools");
  else args.push("--tools", contract.tools.join(","));
  args.push("--mode", contract.channel === "json-result" ? "json" : "rpc");
  if (contract.channel === "json-result") args.push(prompt);
  return args;
}

function pathIsInside(directory: string, candidate: string): boolean {
  const child = relative(directory, candidate);
  return child.length > 0
    && child !== ".."
    && !child.startsWith(`..${sep}`)
    && !isAbsolute(child);
}

function childEnvironment(
  inherited: Readonly<Record<string, string | undefined>>,
  launchContractPath: string,
): Readonly<Record<string, string>> {
  const env: Record<string, string> = {};
  for (const [key, value] of Object.entries(inherited)) {
    if (value !== undefined
      && !isReceiverAuthorityEnvironmentKey(key)
      && !key.startsWith("PI_AMQ_")) env[key] = value;
  }
  env["PI_AMQ_WORKER_LAUNCH_CONTRACT"] = launchContractPath;
  env["PI_AMQ_PARENT_LIVENESS_FD"] = "3";
  env["PI_AMQ_PROTOCOL_FD"] = "4";
  return env;
}

function canonicalStrings(values: readonly string[]): string {
  return JSON.stringify([...values].sort());
}

function strictObject(
  value: unknown,
  allowedKeys: readonly string[],
  name: string,
): Record<string, unknown> {
  if (typeof value !== "object" || value === null || Array.isArray(value)) {
    throw new WorkerLaunchValidationError(`${name} must be an object`);
  }
  const object = value as Record<string, unknown>;
  const unknown = Object.keys(object).filter((key) => !allowedKeys.includes(key));
  if (unknown.length > 0) {
    throw new WorkerLaunchValidationError(`${name} contains unknown fields: ${unknown.join(", ")}`);
  }
  return object;
}

function requiredString(value: unknown, name: string): string {
  if (typeof value !== "string" || value.length === 0) {
    throw new WorkerLaunchValidationError(`${name} must be a non-empty string`);
  }
  return value;
}

function prefixedString(value: unknown, prefix: string, name: string): string {
  const parsed = requiredString(value, name);
  if (!parsed.startsWith(`${prefix}_`)) {
    throw new WorkerLaunchValidationError(`${name} must start with ${prefix}_`);
  }
  return parsed;
}

function exactString<T extends string>(
  value: unknown,
  allowed: readonly T[],
  name: string,
): T {
  const parsed = requiredString(value, name);
  if (!allowed.includes(parsed as T)) throw new WorkerLaunchValidationError(`${name} is invalid`);
  return parsed as T;
}

function stringArray(value: unknown, name: string): readonly string[] {
  if (!Array.isArray(value)) throw new WorkerLaunchValidationError(`${name} must be an array`);
  return value.map((item) => requiredString(item, `${name} entry`));
}

function positiveIntegerValue(value: unknown, name: string): number {
  if (typeof value !== "number" || !Number.isSafeInteger(value) || value <= 0) {
    throw new WorkerLaunchValidationError(`${name} must be a positive integer`);
  }
  return value;
}

function positiveFiniteNumber(value: unknown, name: string): number {
  if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) {
    throw new WorkerLaunchValidationError(`${name} must be positive`);
  }
  return value;
}

function absolutePath(value: string, name: string): string {
  if (!isAbsolute(value)) throw new WorkerLaunchValidationError(`${name} must be absolute`);
  return normalize(value);
}

function requireNonEmpty(value: string, name: string): string {
  if (value.length === 0) throw new WorkerLaunchValidationError(`${name} is required`);
  return value;
}

function positiveInteger(value: number, name: string): number {
  if (!Number.isSafeInteger(value) || value <= 0) {
    throw new WorkerLaunchValidationError(`${name} must be a positive integer`);
  }
  return value;
}