import { createHash } from "node:crypto";
import type { AssistantMessage } from "@earendil-works/pi-ai";
import type {
  AgentStepExecutor,
  AgentStepRequest,
  AgentStepSubmission,
  ConversationRange,
} from "../workflows/types.js";

export type SubmissionResult =
  | { accepted: true; message: string }
  | { accepted: false; message: string };

export type PromptDeliveryKind = "step" | "reminder" | "resume";

export type PromptDelivery = {
  prompt: string;
  contract: AgentStepRequest["contract"];
  presentation?: AgentStepRequest["presentation"];
  kind: PromptDeliveryKind;
  /** True when the agent is known to be mid-run, so delivery must be queued. */
  streaming: boolean;
};

/**
 * Bracketing hooks for conversation linkage: a mark is taken when the step
 * prompt is first delivered, and the recorded entry range since that mark is
 * attached to the accepted submission.
 */
export type ConversationHooks = {
  beginAttempt?: (contract: AgentStepRequest["contract"]) => void;
  recoverAssistant?: (contract: AgentStepRequest["contract"]) => AgentStepSubmission | undefined;
  mark: () => number;
  rangeSince: (mark: number) => ConversationRange | undefined;
};

export type ConversationStepExecutorOptions = {
  /** Deliver a prompt into the pi conversation. */
  sendPrompt: (delivery: PromptDelivery) => void;
  /** Reminders sent when the agent settles without submitting. Default 2. */
  maxNudges?: number;
  /** Conversation linkage hooks, wired to the session recorder. */
  conversation?: ConversationHooks;
  /** Called when the engine aborts a pending agent step. */
  onAbort?: (contract: AgentStepRequest["contract"], reason: unknown) => void;
};

type PendingStep = {
  request: AgentStepRequest;
  resolve: (submission: AgentStepSubmission) => void;
  reject: (error: unknown) => void;
  nudgesSent: number;
  /** Conversation mark taken when the prompt was first delivered. */
  mark: number | null;
  /** Latest finalized assistant message observed for this attempt. */
  assistantMessage?: AssistantMessage;
  cleanup: () => void;
  /** Resolves when this step stops being the pending step. */
  cleared: Promise<void>;
  markCleared: () => void;
};

const DEFAULT_MAX_NUDGES = 2;

/**
 * AgentStepExecutor that runs steps inside the current pi conversation. The
 * engine hands it a prompt; it delivers the prompt as a model-facing workflow
 * message and resolves once the model submits an accepted output through the `workflow`
 * tool. If the agent settles without submitting, it nudges the model a
 * bounded number of times before failing the step.
 */
export class ConversationStepExecutor implements AgentStepExecutor {
  readonly assistantMessageMode = "visible" as const;
  private readonly sendPrompt: (delivery: PromptDelivery) => void;
  private readonly maxNudges: number;
  private readonly conversation: ConversationHooks | undefined;
  private readonly onAbort: ConversationStepExecutorOptions["onAbort"];
  private pending: PendingStep | null = null;
  private streaming = false;
  private heldByUser = false;

  constructor(options: ConversationStepExecutorOptions) {
    this.sendPrompt = options.sendPrompt;
    this.maxNudges = options.maxNudges ?? DEFAULT_MAX_NUDGES;
    this.conversation = options.conversation;
    this.onAbort = options.onAbort;
  }

  /** Track agent streaming state (wire to agent_start / agent_settled). */
  setStreaming(streaming: boolean): void {
    this.streaming = streaming;
  }

  get pendingStepId(): string | null {
    return this.pending?.request.contract.nodeId ?? null;
  }

  get pendingCompletion(): "submit" | "assistant" | null {
    return this.pending === null ? null : completionKind(this.pending.request);
  }

  /**
   * Hold the pending step for the user: no nudges are sent while held, so an
   * escape-interrupted conversation stays quiet until the user resumes.
   */
  hold(): void {
    this.heldByUser = true;
  }

  get held(): boolean {
    return this.heldByUser;
  }

  /**
   * Release a user hold. When a step is still pending, its prompt is
   * re-delivered so the model picks the step back up.
   */
  release(): void {
    if (!this.heldByUser) {
      return;
    }
    this.heldByUser = false;
    const pending = this.pending;
    if (!pending) {
      return;
    }
    try {
      this.conversation?.beginAttempt?.(pending.request.contract);
      this.sendPrompt(this.delivery(pending.request, pending.request.prompt, "resume"));
    } catch (error) {
      this.clearPending();
      pending.reject(error);
    }
  }

  async runAgentStep(request: AgentStepRequest, signal: AbortSignal): Promise<AgentStepSubmission> {
    if (this.pending) {
      throw new Error("Another workflow step is already awaiting output");
    }
    if (completionKind(request) === "assistant") {
      const recovered = this.conversation?.recoverAssistant?.(request.contract);
      if (recovered !== undefined) {
        return recovered;
      }
    }
    return await new Promise<AgentStepSubmission>((resolve, reject) => {
      const onAbort = () => {
        const reason: unknown = signal.reason ?? new Error("Workflow step aborted");
        if (this.pending?.request !== request) {
          return;
        }
        this.clearPending();
        try {
          this.onAbort?.(request.contract, reason);
        } finally {
          reject(reason);
        }
      };
      signal.addEventListener("abort", onAbort, { once: true });
      let markCleared!: () => void;
      const cleared = new Promise<void>((resolveCleared) => {
        markCleared = resolveCleared;
      });
      this.conversation?.beginAttempt?.(request.contract);
      this.pending = {
        request,
        resolve,
        reject,
        nudgesSent: 0,
        mark: this.conversation?.mark() ?? null,
        cleanup: () => signal.removeEventListener("abort", onAbort),
        cleared,
        markCleared,
      };
      if (signal.aborted) {
        onAbort();
        return;
      }
      try {
        this.sendPrompt(this.delivery(request, request.prompt, "step"));
      } catch (error) {
        // A failed delivery must not leave the step installed, or every
        // subsequent agent node would fail with "already awaiting output".
        this.clearPending();
        reject(error);
      }
    });
  }

  /** Called by the `workflow` tool when the model submits a step output. */
  async submit(stepId: string, attemptId: string, output: unknown): Promise<SubmissionResult> {
    const pending = this.pending;
    if (!pending) {
      return {
        accepted: false,
        message:
          "No workflow step is awaiting output. Do not call the workflow tool outside an active workflow step.",
      };
    }
    const expected = pending.request.contract.nodeId;
    if (stepId !== expected) {
      return {
        accepted: false,
        message: `Wrong step id ${JSON.stringify(stepId)}; the pending step is ${JSON.stringify(expected)}.`,
      };
    }
    // Loops revisit the same node id, so a delayed duplicate submission from
    // an earlier attempt would otherwise be accepted as this attempt's output.
    const expectedAttempt = pending.request.contract.attemptId;
    if (attemptId !== expectedAttempt) {
      return {
        accepted: false,
        message: `Stale attempt id ${JSON.stringify(attemptId)} for step ${JSON.stringify(
          stepId,
        )}; the pending attempt is ${JSON.stringify(expectedAttempt)}. Use the attempt id from the latest step contract.`,
      };
    }
    if (completionKind(pending.request) === "assistant") {
      return {
        accepted: false,
        message:
          "This step completes with a normal assistant response. Do not submit workflow output.",
      };
    }
    // Race validation against the step being cleared: a hung `validate`
    // callback must not leave this tool call (and therefore pi) blocked after
    // a timeout or cancel already resolved the run.
    const result = await Promise.race([
      pending.request.accept(output),
      pending.cleared.then(() => null),
    ]);
    // The step may have timed out or been cancelled (and a newer step
    // installed) while validation was awaited; a stale submission must not
    // clear or resolve the newer pending step.
    if (result === null || this.pending !== pending) {
      return {
        accepted: false,
        message: `Step ${JSON.stringify(stepId)} is no longer awaiting output.`,
      };
    }
    if (!result.ok) {
      return {
        accepted: false,
        message: `Output rejected for step ${JSON.stringify(stepId)}: ${result.error}`,
      };
    }
    this.clearPending();
    const conversation =
      pending.mark !== null ? this.conversation?.rangeSince(pending.mark) : undefined;
    pending.resolve({
      output: result.value,
      ...(conversation !== undefined ? { conversation } : {}),
    });
    return {
      accepted: true,
      message: `Output accepted for step ${JSON.stringify(stepId)}.`,
    };
  }

  /** Keep the latest finalized assistant message for an assistant-output step. */
  handleMessageEnd(message: unknown): void {
    const pending = this.pending;
    if (pending === null || completionKind(pending.request) !== "assistant") {
      return;
    }
    const assistant = assistantMessageLike(message);
    if (assistant !== undefined) {
      pending.assistantMessage = assistant as AssistantMessage;
    }
  }

  /**
   * Called when the agent settles. Returns true when a nudge was sent, false
   * when there was nothing to do. Submitted steps use bounded nudges;
   * assistant-output steps accept or fail the one visible response.
   */
  handleAgentSettled(): boolean {
    const pending = this.pending;
    if (!pending) {
      return false;
    }
    if (this.heldByUser) {
      // The user interrupted deliberately; reminding the model now would
      // steal the conversation back. The step waits for an explicit resume.
      return false;
    }
    if (completionKind(pending.request) === "assistant") {
      try {
        const output = visibleAssistantText(
          pending.assistantMessage,
          pending.request.contract.maxOutputChars,
        );
        const conversation =
          pending.mark !== null ? this.conversation?.rangeSince(pending.mark) : undefined;
        const assistantMessage = {
          sha256: createHash("sha256").update(output).digest("hex"),
          ...(conversation?.lastEntryId !== undefined ? { entryId: conversation.lastEntryId } : {}),
          ...(pending.request.contract.maxOutputChars !== undefined
            ? { maxChars: pending.request.contract.maxOutputChars }
            : {}),
        };
        this.clearPending();
        pending.resolve({
          output,
          assistantMessage,
          ...(conversation !== undefined ? { conversation } : {}),
        });
      } catch (error) {
        this.clearPending();
        pending.reject(error);
      }
      return false;
    }
    if (pending.nudgesSent >= this.maxNudges) {
      this.clearPending();
      pending.reject(
        new Error(
          `Agent settled ${pending.nudgesSent + 1} times without submitting step ${JSON.stringify(
            pending.request.contract.nodeId,
          )} via the workflow tool`,
        ),
      );
      return false;
    }
    pending.nudgesSent += 1;
    const { nodeId, attemptId } = pending.request.contract;
    try {
      this.conversation?.beginAttempt?.(pending.request.contract);
      this.sendPrompt(
        this.delivery(
          pending.request,
          [
            `Reminder: workflow step ${JSON.stringify(nodeId)} is still awaiting your output.`,
            "Complete it by calling the `workflow` tool with:",
            `{"action": "submit", "step": ${JSON.stringify(nodeId)}, "attempt": ${JSON.stringify(attemptId)}, "output": <your result>}`,
            `Expected output: ${pending.request.contract.expectedOutput ?? "a JSON object with your result"}`,
          ].join("\n"),
          "reminder",
        ),
      );
    } catch (error) {
      // No reminder turn was started, so nothing would settle the step; fail
      // it promptly instead of waiting out the node timeout.
      this.clearPending();
      pending.reject(error);
      return false;
    }
    return true;
  }

  private delivery(
    request: AgentStepRequest,
    prompt: string,
    kind: PromptDeliveryKind,
  ): PromptDelivery {
    return {
      prompt,
      contract: request.contract,
      ...(request.presentation !== undefined ? { presentation: request.presentation } : {}),
      kind,
      streaming: this.streaming,
    };
  }

  private clearPending(): void {
    this.pending?.cleanup();
    this.pending?.markCleared();
    this.pending = null;
  }
}

function completionKind(request: AgentStepRequest): "submit" | "assistant" {
  return request.contract.completion;
}

function assistantMessageLike(
  message: unknown,
): Pick<AssistantMessage, "role" | "content" | "stopReason" | "errorMessage"> | undefined {
  if (message === null || typeof message !== "object") return undefined;
  const candidate = message as Partial<AssistantMessage>;
  if (candidate.role !== "assistant" || !Array.isArray(candidate.content)) return undefined;
  if (typeof candidate.stopReason !== "string") return undefined;
  return candidate as Pick<AssistantMessage, "role" | "content" | "stopReason" | "errorMessage">;
}

/** Extract the exact visible text blocks from one finalized assistant message. */
export function visibleAssistantText(message: unknown, maxChars?: number): string {
  const assistant = assistantMessageLike(message);
  if (assistant === undefined) {
    throw new Error("Assistant step settled without a final assistant message");
  }
  if (assistant.stopReason !== "stop" && assistant.stopReason !== "length") {
    throw new Error(
      assistant.errorMessage?.trim() ||
        `Assistant step stopped with ${JSON.stringify(assistant.stopReason)} before a final response`,
    );
  }
  const text = assistant.content
    .filter((part) => part.type === "text")
    .map((part) => part.text)
    .join("\n");
  if (text.trim().length === 0) {
    throw new Error("Assistant step returned no visible text");
  }
  if (maxChars !== undefined && text.length > maxChars) {
    throw new Error(
      `Assistant response has ${text.length} characters, above the configured limit of ${maxChars}`,
    );
  }
  return text;
}