import type { Context, JsonValue } from "@earendil-works/chord";
import {
  defineDoc,
  defineEntry,
  defineTask,
  type Harness,
  type TaskRuntime,
} from "@earendil-works/pi-durable";
import type { ShellExecution, ShellState } from "@berry/api/shell.ts";

export const ShellDoc = defineDoc<ShellState>({
  kind: "berry.shell",
  version: 1,
  scope: "conversation",
  history: "latest",
  fork: "initial",
  initial: () => ({}),
  checkpointWhen: () => true,
});
export const ShellEntry = defineEntry<ShellExecution>("berry.shell");

type Input = { command: string; excludeFromContext: boolean };
type State = { phase: "run" };
type Runtime = TaskRuntime<Input, State, null, object>;

/** Passive shell execution: its result is a write, never an input that starts a model run. */
export function shellTask(
  execute: (
    command: string,
    runtime: Runtime,
    context: Context,
    progress: (details: JsonValue) => void,
  ) => Promise<{ value: JsonValue | undefined; text: string }>,
  harness: () => Harness,
) {
  async function finish(runtime: Runtime, context: Context, run: ShellExecution, text: string) {
    const conversation = (await harness().conversation(runtime.conversationId, context))!;
    await conversation.submit({
      type: "write",
      requestId: `shell:${runtime.taskId}`,
      entry: {
        kind: ShellEntry.kind,
        data: run,
        ...(run.excludeFromContext ? {} : {
          model: [{
            role: "user" as const,
            content: `Ran \`${run.command}\`\n\n${text}`,
            timestamp: run.startedAt,
          }],
        }),
      },
    }, context);
    await runtime.commit(async (tx) => {
      delete (await tx.doc(ShellDoc, runtime.conversationId)).run;
      return {
        status: "terminal",
        outcome: run.status === "cancelled"
          ? { status: "aborted" }
          : { status: "completed", result: null },
      };
    }, context);
  }

  return defineTask<Input, State, null>({
    name: "berry.shell",
    version: 1,
    initial: () => ({ phase: "run" }),
    phases: {
      async run(_, runtime, context) {
        const run = structuredClone(
          (await runtime.snapshot(ShellDoc, runtime.conversationId, context))?.run,
        );
        if (!run) throw new Error("Shell execution has no live state");
        // shell side effects cannot safely be replayed after a restart
        const invocation = crypto.randomUUID();
        const started = await runtime.memo("started", invocation, context);
        let pending = Promise.resolve();
        let text = "";
        try {
          if (started !== invocation) {
            throw new Error("Command interrupted by a restart; not rerun.");
          }
          const result = await execute(run.command, runtime, context, (details) => {
            run.details = details;
            pending = pending.then(() =>
              runtime.commit(async (tx) => {
                const live = (await tx.doc(ShellDoc, runtime.conversationId)).run;
                if (live) live.details = details;
              }, context)
            );
            void pending.catch(() => {});
          });
          run.result = result.value ?? null;
          run.status = "ok";
          text = result.text;
        } catch (error) {
          if (runtime.signal.aborted) throw error;
          run.status = "error";
          run.error = error instanceof Error ? error.message : String(error);
          text = `${output(run)}\n\n${run.error}`;
        }
        await pending;
        run.durationMs = runtime.now() - run.startedAt;
        await finish(runtime, context, run, text);
      },
    },
    async abort(_, runtime, context) {
      const run = structuredClone(
        (await runtime.snapshot(ShellDoc, runtime.conversationId, context))?.run,
      );
      if (!run) {
        await runtime.commit(
          () => ({ status: "terminal", outcome: { status: "aborted" } }),
          context,
        );
        return;
      }
      run.status = "cancelled";
      run.error = "Command cancelled.";
      run.durationMs = runtime.now() - run.startedAt;
      await finish(runtime, context, run, `${output(run)}\n\n(command cancelled)`);
    },
  });
}

function output(run: ShellExecution): string {
  const details = run.details as { rawOutput?: string } | null;
  return details?.rawOutput ?? "(no output)";
}