import { tool } from "ai";
import * as Stream from "effect/Stream";
import { z } from "zod";
import {
  SHELL_INPUT_BYTES,
  SHELL_MAX_MS,
  SHELL_OUTPUT_BYTES,
  type ShellResult,
} from "../shared/shell";
import { shellResults } from "./shell-execution";
import type { ToolActivityDescriptor } from "./tool-activity";

export const shellInput = z
  .object({
    command: z
      .string()
      .min(1)
      .max(SHELL_INPUT_BYTES)
      .refine(
        (value) =>
          new TextEncoder().encode(value).byteLength <= SHELL_INPUT_BYTES,
        "Command must fit in 32 KiB",
      )
      .describe(
        "A complete Bash command. For quote-heavy or multiline JavaScript, use a single-quoted heredoc such as node <<'JS' followed by the script and JS on its own line; do not wrap the script in nested node -e quotes.",
      ),
    timeoutMs: z.number().int().min(1000).max(SHELL_MAX_MS).default(30_000),
  })
  .strict();

export interface ShellHost {
  reserve(): Promise<{ id: string; expiresAt: number }>;
  launch(id: string, command: string): Promise<ReadableStream<Uint8Array>>;
  close(id: string): Promise<boolean>;
}

/** Ordinary native AI SDK preliminary results; Think owns transport/history. */
export function createShellTool(host: (timeoutMs: number) => ShellHost) {
  return tool({
    description:
      "Run a Bash command/script in a fresh temporary Linux sandbox with Node.js and Bun. Use heredocs to write and process temporary files in this invocation. No files persist between calls. Streams bounded stdout/stderr; kills the container at completion, cancellation or limit. No Python runtime is promised.",
    inputSchema: shellInput,
    // Think preserves preliminary output only for async-generator execute functions.
    execute: async function* ({ command, timeoutMs }, { abortSignal }) {
      yield* Stream.toAsyncIterable(
        shellResults(() => host(timeoutMs), command, abortSignal),
      );
    },
  });
}

export const shellActivity: ToolActivityDescriptor = {
  kind: "shell",
  label: "Run temporary shell",
  outcome: (output) =>
    (output as ShellResult).status === "succeeded" ? "succeeded" : "failed",
  outputSummary: (output) =>
    (output as ShellResult).cleanup === "pending"
      ? "Command completed; workspace cleanup pending"
      : "Temporary shell completed",
  progress: (output) => {
    const value = output as ShellResult;
    return Number.isSafeInteger(value?.outputBytes)
      ? {
          text: "Receiving shell output",
          completed: value.outputBytes,
          total: SHELL_OUTPUT_BYTES,
        }
      : undefined;
  },
};