import type {
  AgentToolResult,
  ToolDefinition,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import {
  type EvalDetachedCellManager,
  type EvalDetachedCellSnapshot,
  type EvalDetachedCellState,
  isTerminalDetachedState,
} from "./detached-cell-manager.ts";

export const WAIT_DEFAULT_YIELD_TIME_MS = 10_000;
export const WAIT_DEFAULT_MAX_TOKENS = 10_000;
export const TOKENS_TO_CHARS = 4;

export interface WaitToolParams {
  readonly cell_id: string;
  readonly max_tokens?: number;
  readonly terminate?: boolean;
  readonly yield_time_ms?: number;
}

export interface WaitToolDetails {
  readonly cellId?: string;
  readonly isError?: boolean;
  readonly newText?: string;
  readonly state?: EvalDetachedCellState;
}

const waitParameters = Type.Object(
  {
    cell_id: Type.String({ minLength: 1 }),
    yield_time_ms: Type.Optional(Type.Number({ minimum: 1 })),
    max_tokens: Type.Optional(Type.Number({ minimum: 1 })),
    terminate: Type.Optional(Type.Boolean()),
  },
  { additionalProperties: false }
);

const WAIT_DESCRIPTION = [
  "Use when a detached eval cell is still running: poll its new output with the cell_id from the eval result.",
  "wait resumes a detached eval cell.",
  "cell_id identifies the cell returned by a previous eval call.",
  "The call returns new output since the previous wait.",
  "yield_time_ms sets the wait budget before yielding again.",
  "max_tokens sets the output budget. The budget is approximate: one token equals four characters.",
  "terminate: true stops the cell.",
].join("\n");

const STILL_RUNNING_NOTE =
  "(cell still running; use wait again with this cell_id)";

export function createWaitTool(
  cellManager:
    | EvalDetachedCellManager
    | (() => EvalDetachedCellManager | undefined)
): ToolDefinition<typeof waitParameters, WaitToolDetails, unknown> {
  return {
    name: "wait",
    label: "Wait",
    description: WAIT_DESCRIPTION,
    parameters: waitParameters,
    executionMode: "sequential",
    async execute(
      _toolCallId,
      params,
      _signal,
      _onUpdate,
      _ctx
    ): Promise<AgentToolResult<WaitToolDetails>> {
      try {
        validateParams(params);
        const manager =
          typeof cellManager === "function" ? cellManager() : cellManager;
        if (manager === undefined) {
          throw new Error("No active codemode session");
        }
        if (params.terminate === true) {
          const snapshot = await manager.stop(params.cell_id);
          return waitResult(snapshot, stopText(snapshot));
        }
        const yieldTimeMs = params.yield_time_ms ?? WAIT_DEFAULT_YIELD_TIME_MS;
        const maxTokens = params.max_tokens ?? WAIT_DEFAULT_MAX_TOKENS;
        const { snapshot, newText } = await manager.waitForOutput(
          params.cell_id,
          {
            yieldTimeMs,
            maxChars: maxTokens * TOKENS_TO_CHARS,
          }
        );
        if (isTerminalDetachedState(snapshot.state)) {
          return waitResult(snapshot, terminalText(snapshot, newText));
        }
        const text =
          newText.length === 0
            ? STILL_RUNNING_NOTE
            : `${newText}\n${STILL_RUNNING_NOTE}`;
        return waitResult(snapshot, text);
      } catch (error) {
        return {
          content: [
            {
              type: "text",
              text: error instanceof Error ? error.message : String(error),
            },
          ],
          details: { isError: true },
        };
      }
    },
  };
}

function validateParams(params: WaitToolParams): void {
  if (typeof params.cell_id !== "string" || params.cell_id.length < 1) {
    throw new RangeError('wait requires a non-empty "cell_id"');
  }
  for (const key of Object.keys(params)) {
    if (
      key !== "cell_id" &&
      key !== "yield_time_ms" &&
      key !== "max_tokens" &&
      key !== "terminate"
    ) {
      throw new RangeError(`wait does not accept parameter "${key}"`);
    }
  }
}

function terminalText(
  snapshot: EvalDetachedCellSnapshot,
  newText: string
): string {
  // For the error state the failure message lives in the snapshot's text
  // content (detachedErrorResult appends it there), not in the output tail;
  // surface the full terminal text the way peek/stop do.
  if (snapshot.state !== "error") {
    return newText;
  }
  return snapshot.result.content
    .filter((part) => part.type === "text")
    .map((part) => part.text)
    .join("\n");
}

function stopText(snapshot: EvalDetachedCellSnapshot): string {
  return snapshot.outputTail.length === 0
    ? "(cell stopped)"
    : snapshot.outputTail;
}

function waitResult(
  snapshot: EvalDetachedCellSnapshot,
  text: string
): AgentToolResult<WaitToolDetails> {
  return {
    content: [{ type: "text", text }],
    details: {
      cellId: snapshot.cellId,
      state: snapshot.state,
      newText: text,
      ...(snapshot.state === "error" ? { isError: true } : {}),
    },
  };
}