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

import type { AgentManager } from "./agent-manager.js";
import { getDefaultMaxTurns, normalizeMaxTurns } from "./agent-runner.js";
import { getAgentConfig, resolveType } from "./agent-types.js";
import { resolveAgentInvocationConfig } from "./invocation-config.js";
import { type ModelRegistry, resolveModel } from "./model-resolver.js";
import type { AgentRecord, RunnerBackend, SubagentType } from "./types.js";
import type {
  WorkflowAgentChild,
  WorkflowAgentRequest,
  WorkflowAgentResult,
  WorkflowAgentRunContext,
  WorkflowAgentRunner,
} from "./workflow-types.js";

const DEFAULT_WORKFLOW_AGENT_TYPE = "general-purpose";
const STRUCTURED_OUTPUT_TOOL_NAME = "structured_output";
const STRUCTURED_OUTPUT_CLEANUP_TIMEOUT_MS = 1000;

interface StructuredOutputCapture {
  isCaptured: () => boolean;
  promise: Promise<void>;
}

export interface WorkflowAgentManagerRunnerOptions {
  pi: ExtensionAPI;
  ctx: ExtensionContext;
  manager: Pick<AgentManager, "abort" | "getRecord" | "spawn">;
  runnerBackend?: RunnerBackend;
  reloadAgents?: () => void;
  signal?: AbortSignal;
  onChildUpdate?: (child: WorkflowAgentChild) => void;
}

export function createWorkflowAgentRunner(
  options: WorkflowAgentManagerRunnerOptions
): WorkflowAgentRunner {
  return async (request, context) =>
    runWorkflowAgent(request, context, options);
}

async function runWorkflowAgent(
  request: WorkflowAgentRequest,
  context: WorkflowAgentRunContext,
  options: WorkflowAgentManagerRunnerOptions
): Promise<WorkflowAgentResult> {
  options.reloadAgents?.();

  const rawType = getRequestedAgentType(request);
  const resolvedType = resolveType(rawType);
  const agentConfig = resolvedType ? getAgentConfig(resolvedType) : undefined;
  if (!(resolvedType && agentConfig) || agentConfig.enabled === false) {
    throw new Error(
      `Agent type "${rawType}" is not defined. Create it in .pi/agents/${rawType}.md or ~/.pi/agent/agents/${rawType}.md, then retry.`
    );
  }

  const invocation = resolveAgentInvocationConfig(agentConfig, {
    model: stringField(request.model),
    thinking: stringField(request.thinking),
    max_turns: numberField(request.max_turns),
    context: contextField(request.context),
    isolated: booleanField(request.isolated),
    isolation: isolationField(request.isolation),
  });

  let model = options.ctx.model;
  if (invocation.modelInput) {
    const resolvedModel = resolveModel(
      invocation.modelInput,
      options.ctx.modelRegistry as ModelRegistry
    );
    if (typeof resolvedModel === "string") {
      if (invocation.modelFromParams) {
        throw new Error(resolvedModel);
      }
    } else {
      model = resolvedModel;
    }
  }

  const structuredOutputSchema = getStructuredOutputSchema(request);
  const runnerBackend =
    structuredOutputSchema !== undefined && options.runnerBackend === "herdr"
      ? "in-process"
      : options.runnerBackend;
  const effectiveMaxTurns = normalizeMaxTurns(
    runnerBackend === "herdr"
      ? invocation.maxTurns
      : (invocation.maxTurns ?? getDefaultMaxTurns())
  );

  let structuredOutput: unknown;
  let structuredOutputCaptured = false;
  let resolveStructuredOutputCapture: (() => void) | undefined;
  const structuredOutputCapturePromise = new Promise<void>((resolve) => {
    resolveStructuredOutputCapture = resolve;
  });
  const customTools =
    structuredOutputSchema === undefined
      ? undefined
      : [
          createStructuredOutputTool(structuredOutputSchema, (value) => {
            structuredOutput = value;
            if (!structuredOutputCaptured) {
              structuredOutputCaptured = true;
              resolveStructuredOutputCapture?.();
            }
          }),
        ];

  const signal = context.signal ?? options.signal;
  if (signal?.aborted) {
    throw createWorkflowAgentAbortError();
  }

  const prompt =
    structuredOutputSchema === undefined
      ? request.prompt
      : appendStructuredOutputInstruction(
          request.prompt,
          structuredOutputSchema
        );

  const id = options.manager.spawn(
    options.pi,
    options.ctx,
    resolvedType as SubagentType,
    prompt,
    {
      description: getDescription(request, resolvedType, context),
      model,
      maxTurns: effectiveMaxTurns,
      isolated: invocation.isolated,
      inheritContext: invocation.inheritContext,
      thinkingLevel: invocation.thinking,
      runnerBackend,
      isBackground: true,
      allowAskParent: false,
      isolation: invocation.isolation,
      signal,
      customTools,
    }
  );

  const abortChild = () => options.manager.abort(id);
  signal?.addEventListener("abort", abortChild, { once: true });

  try {
    const initialRecord = options.manager.getRecord(id);
    if (initialRecord) {
      initialRecord.resultConsumed = true;
      emitChildUpdate(options, initialRecord);
    }

    const record = await waitForWorkflowChild(
      options.manager,
      id,
      signal,
      structuredOutputSchema === undefined
        ? undefined
        : {
            isCaptured: () => structuredOutputCaptured,
            promise: structuredOutputCapturePromise,
          }
    );
    const completedFromStructuredOutput = structuredOutputCaptured;
    const resultStatus = completedFromStructuredOutput
      ? "completed"
      : record.status;
    emitChildUpdate(options, record, resultStatus);
    if (
      !(
        completedFromStructuredOutput ||
        isSuccessfulWorkflowChildStatus(record.status)
      )
    ) {
      throw createWorkflowChildFailedError(record);
    }
    if (
      structuredOutputSchema !== undefined &&
      structuredOutput === undefined
    ) {
      throw new Error("Workflow agent structured output was not captured.");
    }

    return {
      id,
      type: record.type,
      status: resultStatus,
      result: record.result,
      error: record.error,
      structuredOutput,
      warnings: record.warnings,
      toolUses: record.toolUses,
    };
  } finally {
    signal?.removeEventListener("abort", abortChild);
  }
}

function getRequestedAgentType(request: WorkflowAgentRequest): string {
  return (
    stringField(request.agent) ??
    stringField(request.subagent_type) ??
    stringField(request.type) ??
    DEFAULT_WORKFLOW_AGENT_TYPE
  );
}

function getStructuredOutputSchema(
  request: WorkflowAgentRequest
): unknown | undefined {
  return request.schema === undefined ? request.output : request.schema;
}

function getDescription(
  request: WorkflowAgentRequest,
  type: string,
  context: WorkflowAgentRunContext
): string {
  if (typeof request.description === "string" && request.description.trim()) {
    return request.description;
  }
  return context.phase
    ? `workflow:${context.phase}:${type}`
    : `workflow:${type}`;
}

function createStructuredOutputTool(
  schema: unknown,
  onOutput: (value: unknown) => void
): ToolDefinition {
  const parameters = isSchemaObject(schema)
    ? schema
    : Type.Object({ value: Type.Unknown() });

  return {
    name: STRUCTURED_OUTPUT_TOOL_NAME,
    label: "Structured Output",
    description:
      "Return the final structured workflow result. Use this as the last action.",
    promptSnippet: "Emit a final structured workflow result",
    promptGuidelines: [
      "Use structured_output as your final action when the workflow request includes an output schema.",
      "After calling structured_output, do not emit another assistant response in the same turn.",
    ],
    parameters: parameters as ToolDefinition["parameters"],
    execute(_toolCallId, params) {
      onOutput(params);
      return Promise.resolve({
        content: [
          { type: "text" as const, text: "Structured output captured." },
        ],
        details: params,
        terminate: true,
      });
    },
  };
}

async function waitForWorkflowChild(
  manager: Pick<AgentManager, "abort" | "getRecord">,
  id: string,
  signal?: AbortSignal,
  structuredOutputCapture?: StructuredOutputCapture
): Promise<AgentRecord> {
  while (true) {
    const record = manager.getRecord(id);
    if (!record) {
      throw new Error(
        `Workflow child agent "${id}" was removed before completion.`
      );
    }
    if (isTerminalStatus(record.status)) {
      return record;
    }
    if (structuredOutputCapture?.isCaptured()) {
      return stopCapturedWorkflowChild(manager, id, record, signal);
    }
    const childWait = waitForPromiseOrAbort<unknown>(
      record.promise ?? delay(20),
      signal
    );
    if (!structuredOutputCapture) {
      await childWait;
      continue;
    }

    const captured = await Promise.race([
      childWait.then(() => false),
      structuredOutputCapture.promise.then(() => true),
    ]);
    if (captured) {
      return stopCapturedWorkflowChild(manager, id, record, signal);
    }
  }
}

async function stopCapturedWorkflowChild(
  manager: Pick<AgentManager, "abort" | "getRecord">,
  id: string,
  record: AgentRecord,
  signal?: AbortSignal
): Promise<AgentRecord> {
  if (!isTerminalStatus(record.status)) {
    manager.abort(id);
  }

  if (record.promise) {
    try {
      await waitForPromiseOrAbortWithTimeout(
        record.promise,
        signal,
        STRUCTURED_OUTPUT_CLEANUP_TIMEOUT_MS
      );
    } catch (error) {
      record.session?.dispose?.();
      throw error;
    }
  }

  const completedRecord = manager.getRecord(id);
  if (!completedRecord) {
    throw new Error(
      `Workflow child agent "${id}" was removed during structured output cleanup.`
    );
  }
  if (!isTerminalStatus(completedRecord.status)) {
    throw new Error(
      `Workflow child agent "${id}" did not stop after structured output capture.`
    );
  }
  return completedRecord;
}

function waitForPromiseOrAbortWithTimeout<T>(
  promise: Promise<T>,
  signal: AbortSignal | undefined,
  timeoutMs: number
): Promise<T> {
  let timeout: ReturnType<typeof setTimeout> | undefined;
  const timeoutPromise = new Promise<never>((_resolve, reject) => {
    timeout = setTimeout(() => {
      reject(
        new Error(
          `Workflow child cleanup exceeded ${timeoutMs}ms after structured output capture.`
        )
      );
    }, timeoutMs);
  });
  return Promise.race([
    waitForPromiseOrAbort(promise, signal),
    timeoutPromise,
  ]).finally(() => {
    if (timeout) {
      clearTimeout(timeout);
    }
  });
}

function waitForPromiseOrAbort<T>(
  promise: Promise<T>,
  signal?: AbortSignal
): Promise<T> {
  if (!signal) {
    return promise;
  }
  if (signal.aborted) {
    return Promise.reject(createWorkflowAgentAbortError());
  }
  return new Promise<T>((resolve, reject) => {
    const cleanup = () => signal.removeEventListener("abort", onAbort);
    const onAbort = () => {
      cleanup();
      reject(createWorkflowAgentAbortError());
    };
    signal.addEventListener("abort", onAbort, { once: true });
    promise.then(
      (value) => {
        cleanup();
        resolve(value);
      },
      (error) => {
        cleanup();
        reject(error);
      }
    );
  });
}

function createWorkflowAgentAbortError(): Error {
  const error = new Error("Workflow agent call aborted.");
  error.name = "AbortError";
  return error;
}

function delay(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

function emitChildUpdate(
  options: WorkflowAgentManagerRunnerOptions,
  record: AgentRecord,
  status = record.status
): void {
  options.onChildUpdate?.({
    id: record.id,
    type: record.type,
    status,
    description: record.description,
  });
}

function createWorkflowChildFailedError(record: AgentRecord): Error {
  const details = record.error ? `: ${record.error}` : "";
  return new Error(
    `Workflow child agent "${record.id}" ended with status "${record.status}"${details}.`
  );
}

function appendStructuredOutputInstruction(
  prompt: string,
  schema: unknown
): string {
  return [
    prompt,
    "",
    `Call ${STRUCTURED_OUTPUT_TOOL_NAME} as your final action with output matching this schema:`,
    JSON.stringify(schema),
  ].join("\n");
}

function isTerminalStatus(status: AgentRecord["status"]): boolean {
  return !["queued", "running"].includes(status);
}

function isSuccessfulWorkflowChildStatus(
  status: AgentRecord["status"]
): boolean {
  return status === "completed" || status === "steered";
}

function isSchemaObject(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null && "type" in value;
}

function stringField(value: unknown): string | undefined {
  return typeof value === "string" ? value : undefined;
}

function numberField(value: unknown): number | undefined {
  return typeof value === "number" ? value : undefined;
}

function booleanField(value: unknown): boolean | undefined {
  return typeof value === "boolean" ? value : undefined;
}

function contextField(value: unknown): "fresh" | "fork" | undefined {
  return value === "fresh" || value === "fork" ? value : undefined;
}

function isolationField(value: unknown): "worktree" | undefined {
  return value === "worktree" ? value : undefined;
}