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

const Difficulty = Type.Union([
  Type.Literal("easy"),
  Type.Literal("medium"),
  Type.Literal("hard"),
]);
const Channel = Type.Union([Type.Literal("json-result"), Type.Literal("rpc")]);
const SideEffectPolicy = Type.Union([
  Type.Literal("none"),
  Type.Literal("idempotent"),
  Type.Literal("reconcilable"),
  Type.Literal("external"),
]);
const WorkerContext = Type.Union([Type.Literal("fresh"), Type.Literal("fork")]);
const ReportField = Type.Object({
  type: Type.Union([
    Type.Literal("string"), Type.Literal("number"), Type.Literal("boolean"),
    Type.Literal("object"), Type.Literal("array"), Type.Literal("null"),
  ]),
});
const ReportSchema = Type.Object({
  schema: Type.Literal(1),
  required: Type.Array(Type.String()),
  properties: Type.Record(Type.String(), ReportField),
});
const ReportSchemaPolicy = Type.Union([
  Type.Literal("strict"),
  Type.Literal("preferred"),
  Type.Literal("advisory"),
]);
const AcceptanceCheck = Type.Object({
  command: Type.String(),
  args: Type.Array(Type.String()),
});
const IntegerBudgetLimit = Type.Union([
  Type.Object({ kind: Type.Literal("limited"), value: Type.Integer({ minimum: 1 }) }),
  Type.Object({ kind: Type.Literal("unlimited") }),
]);
const NumberBudgetLimit = Type.Union([
  Type.Object({ kind: Type.Literal("limited"), value: Type.Number({ exclusiveMinimum: 0 }) }),
  Type.Object({ kind: Type.Literal("unlimited") }),
]);
const WorkerBudget = Type.Object({
  max_attempts: IntegerBudgetLimit,
  max_wall_clock_ms: IntegerBudgetLimit,
  max_tokens: IntegerBudgetLimit,
  max_cost_usd: NumberBudgetLimit,
  max_output_bytes: IntegerBudgetLimit,
});
const Model = Type.Object({
  provider: Type.String(),
  id: Type.String(),
  thinking_level: Type.Optional(Type.String()),
});
const Workspace = Type.Object({
  path: Type.String(),
  mode: Type.Union([Type.Literal("read-only"), Type.Literal("mutating")]),
  worktree_path: Type.Optional(Type.String()),
});
const ContentRef = Type.Object({
  kind: Type.Union([Type.Literal("file"), Type.Literal("session"), Type.Literal("artifact")]),
  value: Type.String(),
  hash: Type.Optional(Type.String()),
});

const SpawnParameters = Type.Object({
  request_key: Type.String(),
  difficulty: Type.Optional(Difficulty),
  role: Type.Optional(Type.String({ minLength: 1 })),
  audits_worker_id: Type.Optional(Type.String()),
  continue_worker_id: Type.Optional(Type.String()),
  context: Type.Optional(WorkerContext),
  goal: Type.String(),
  acceptance: Type.Array(Type.String(), { minItems: 1 }),
  constraints: Type.Array(Type.String()),
  budget: Type.Optional(WorkerBudget),
  model: Type.Optional(Model),
  fallback_models: Type.Optional(Type.Array(Model, { maxItems: 4 })),
  capabilities: Type.Optional(Type.Array(Type.String())),
  workspace: Workspace,
  context_refs: Type.Array(ContentRef),
  channel: Channel,
  side_effect_policy: SideEffectPolicy,
  report_schema: Type.Optional(ReportSchema),
  report_schema_policy: Type.Optional(ReportSchemaPolicy),
  requires_human_approval: Type.Optional(Type.Boolean()),
  acceptance_check: Type.Optional(AcceptanceCheck),
});
const ListParameters = Type.Object({
  limit: Type.Optional(Type.Integer({ minimum: 1, maximum: 100 })),
});
const CollectParameters = Type.Object({
  worker_ids: Type.Array(Type.String(), { minItems: 1, maxItems: 32 }),
  structured_reports: Type.Optional(Type.Record(Type.String(), Type.Unknown())),
});
const CancelParameters = Type.Object({ worker_id: Type.String(), reason: Type.String() });
const SteerParameters = Type.Object({
  worker_id: Type.String(),
  message: Type.String(),
  correlation: Type.Optional(Type.String({ minLength: 1, maxLength: 256 })),
});
const ReviseParameters = Type.Object({
  worker_id: Type.String(),
  correlation: Type.Optional(Type.String({ minLength: 1, maxLength: 256 })),
  goal: Type.Optional(Type.String()),
  acceptance: Type.Optional(Type.Array(Type.String(), { minItems: 1 })),
  constraints: Type.Optional(Type.Array(Type.String())),
  report_schema: Type.Optional(ReportSchema),
  report_schema_policy: Type.Optional(ReportSchemaPolicy),
  requires_human_approval: Type.Optional(Type.Boolean()),
  acceptance_check: Type.Optional(AcceptanceCheck),
});
const CleanupRunsParameters = Type.Object({
  older_than_ms: Type.Integer({ minimum: 1 }),
  dry_run: Type.Optional(Type.Boolean()),
  yes: Type.Optional(Type.Boolean()),
});
const WorkflowDoctorParameters = Type.Object({});

export type SpawnWorkerToolParameters = Static<typeof SpawnParameters>;
export type ListWorkersToolParameters = Static<typeof ListParameters>;
export type CollectWorkerToolParameters = Static<typeof CollectParameters>;
export type CancelWorkerToolParameters = Static<typeof CancelParameters>;
export type SteerWorkerToolParameters = Static<typeof SteerParameters>;
export type ReviseContractToolParameters = Static<typeof ReviseParameters>;
export type CleanupRunsToolParameters = Static<typeof CleanupRunsParameters>;

export interface WorkerToolProvider {
  spawn(
    params: SpawnWorkerToolParameters,
    signal: AbortSignal | undefined,
    ctx: ExtensionContext,
  ): Promise<unknown>;
  list(params: ListWorkersToolParameters, ctx: ExtensionContext): Promise<unknown>;
  collect(
    params: CollectWorkerToolParameters,
    signal: AbortSignal | undefined,
    ctx: ExtensionContext,
  ): Promise<unknown>;
  cancel(
    params: CancelWorkerToolParameters,
    signal: AbortSignal | undefined,
    ctx: ExtensionContext,
  ): Promise<unknown>;
  steer(
    params: SteerWorkerToolParameters,
    signal: AbortSignal | undefined,
    ctx: ExtensionContext,
  ): Promise<unknown>;
  reviseContract(
    params: ReviseContractToolParameters,
    signal: AbortSignal | undefined,
    ctx: ExtensionContext,
  ): Promise<unknown>;
  cleanupRuns(
    params: CleanupRunsToolParameters,
    signal: AbortSignal | undefined,
    ctx: ExtensionContext,
  ): Promise<unknown>;
  doctorWorkflow(signal: AbortSignal | undefined, ctx: ExtensionContext): Promise<unknown>;
}

export function registerWorkerTools(pi: ExtensionAPI, provider: WorkerToolProvider): void {
  pi.registerTool({
    name: "spawn_worker",
    label: "Spawn Worker",
    description: "Create one durable worker contract and start its supervised process, optionally continuing a proven-dead terminal worker session.",
    promptSnippet: "Spawn a bounded supervised worker",
    promptGuidelines: [
      "Choose scout, worker, reviewer, oracle, or researcher; user role declarations may override these builtins.",
      "Worker context defaults to fresh; context=fork copies the driver current leaf into an independent child session without changing the parent.",
      "fallback_models is an ordered bounded chain used only for child launch-validation or provider auth, quota, and capacity failures.",
      "Explicit difficulty and capabilities override role defaults. A reviewer auditing a named worker is floored to that worker's model capability.",
      "When tier selection is unavailable or uncertain, omit model so the current model is used.",
      "Keep acceptance criteria falsifiable and grant only the required workspace and capabilities.",
      "budget.max_output_bytes counts every byte Pi writes on its event stream, not the answer size: a worker that only says hello already emits several kilobytes, so budget megabytes rather than kilobytes.",
      "budget.max_tokens counts Pi usage.totalTokens, including the worker input context, not just its visible answer; a trivial worker can start above 2,000 tokens.",
      "budget.max_wall_clock_ms may not exceed the run's total wall-clock allowance, which a refusal reports. The default follows the remaining run deadline.",
      "Default token, cost, and output limits reserve the maximum remaining additive run capacity for the worker. Omit budget for ordinary tasks; declare a tighter share only when the user or task genuinely requires it.",
      "Never invent a smaller budget from apparent task size. Unless the task requires a different hard share, omit budget and keep the runtime maximum defaults.",
      "Use report_schema_policy strict only when a machine consumer requires exact structured fields; preferred (default) keeps raw AMQ/candidate output usable when formatting is wrong; advisory treats the schema as a prompt hint only.",
      "Use continue_worker_id only for deliberate continuation after the prior worker is terminal and proven dead; the new request remains a fresh contract.",
      "If the user does not specify a budget, omit budget; the runtime chooses the maximum run-safe default.",
    ],
    parameters: SpawnParameters,
    async execute(_id, params, signal, _onUpdate, ctx) {
      return invoke(() => provider.spawn(params, signal, ctx));
    },
  });

  pi.registerTool({
    name: "list_workers",
    label: "List Workers",
    description: "List the current run's bounded durable worker, Gate, revision, and stall projections.",
    parameters: ListParameters,
    async execute(_id, params, _signal, _onUpdate, ctx) {
      return invoke(() => provider.list(params, ctx));
    },
  });

  pi.registerTool({
    name: "collect",
    label: "Collect Workers",
    description: "Observe worker candidates and Gate facts without treating collection itself as done.",
    parameters: CollectParameters,
    async execute(_id, params, signal, _onUpdate, ctx) {
      return invoke(() => provider.collect(params, signal, ctx));
    },
  });

  pi.registerTool({
    name: "cancel",
    label: "Cancel Worker",
    description: "Record the sole cancellation verb before signaling the supervised process group.",
    parameters: CancelParameters,
    async execute(_id, params, signal, _onUpdate, ctx) {
      return invoke(() => provider.cancel(params, signal, ctx));
    },
  });

  pi.registerTool({
    name: "steer",
    label: "Steer Worker",
    description: "Deliver already-authorized new input to an RPC worker; waking a typed wait requires its exact correlation.",
    parameters: SteerParameters,
    async execute(_id, params, signal, _onUpdate, ctx) {
      return invoke(() => provider.steer(params, signal, ctx));
    },
  });

  pi.registerTool({
    name: "revise_contract",
    label: "Revise Worker Contract",
    description: "Record an immutable contract revision and deliver it at the next inference boundary; waking a typed wait requires its exact correlation.",
    parameters: ReviseParameters,
    async execute(_id, params, signal, _onUpdate, ctx) {
      return invoke(() => provider.reviseContract(params, signal, ctx));
    },
  });

  pi.registerTool({
    name: "cleanup_runs",
    label: "Cleanup Workflow Runs",
    description: "Preview or explicitly delete old terminal workflow runs, worker sessions, and artifacts.",
    promptGuidelines: [
      "Dry-run is the default. Set dry_run=false only when deletion is intended.",
      "TUI deletion always requires user confirmation; headless deletion also requires yes=true.",
    ],
    parameters: CleanupRunsParameters,
    async execute(_id, params, signal, _onUpdate, ctx) {
      return invoke(() => provider.cleanupRuns(params, signal, ctx));
    },
  });

  pi.registerTool({
    name: "workflow_doctor",
    label: "Workflow Doctor",
    description: "Read-only workflow diagnostics for run ledgers, processes, leases, sessions, and disk usage.",
    parameters: WorkflowDoctorParameters,
    async execute(_id, _params, signal, _onUpdate, ctx) {
      return invoke(() => provider.doctorWorkflow(signal, ctx));
    },
  });
}

async function invoke(operation: () => Promise<unknown>): Promise<AgentToolResult<unknown>> {
  try {
    const details = await operation();
    return rendered(details);
  } catch (error) {
    return rendered({ error: { category: "worker-runtime", message: toError(error).message } });
  }
}

function rendered(details: unknown): AgentToolResult<unknown> {
  return {
    content: [{ type: "text", text: JSON.stringify(details, null, 2) }],
    details,
  };
}

function toError(error: unknown): Error {
  return error instanceof Error ? error : new Error(String(error));
}