import {
  Array as EffectArray,
  Chunk,
  Effect,
  Layer,
  Option,
  Runtime,
  Stream,
  Inspectable,
  Schema,
} from "effect";

import { SCRATCH_CELL_ID } from "../constants.ts";
import { scratchCellNotificationsToVsCodeOutput } from "../kernel/CellExecutions.ts";
import { NotebookRuntime } from "../kernel/NotebookRuntime.ts";
import { signalFromToken } from "../lib/signalFromToken.ts";
import { VsCode } from "../platform/VsCode.ts";
import {
  extractCellIdFromCellMessage,
  MarimoNotebookDocument,
} from "../schemas/MarimoNotebookDocument.ts";
import type { CellOperationNotification } from "../types.ts";

/**
 * The agent-facing tool name. Mirrors the `contributes.languageModelTools`
 * entry in package.json — VS Code requires both to match.
 */
const EXECUTE_CODE_TOOL = "marimo_executeCode";

/**
 * Extract a cell-op's stdout/stderr text
 *
 * Read the notification's `console` delta and concatenate each stdout/stderr
 * output's `data`. Other channels (stdin, pdb, media) are skipped, matching
 * marimo's /api/kernel/execute SSE endpoint.
 */
export function consoleText(op: CellOperationNotification): string {
  if (op.console == null) {
    return "";
  }
  return EffectArray.ensure(op.console)
    .filter(
      (output) =>
        output != null &&
        (output.channel === "stdout" || output.channel === "stderr"),
    )
    .map((output) =>
      typeof output.data === "string"
        ? output.data
        : Inspectable.format(output.data),
    )
    .join("");
}

export const ExecuteCodeInput = Schema.Struct({
  /** URI of the marimo notebook whose kernel to run in (explicit, no default). */
  notebookUri: Schema.String,
  /** Python to run in that notebook's live kernel (scratchpad). */
  code: Schema.String,
});
/**
 * Registers the `execute_code` Language Model Tool: the single channel an agent
 * uses to run Python in a marimo notebook's kernel. Exploration stays in the
 * scratchpad; durable edits happen when the agent's code uses
 * `marimo._code_mode` (taught by the marimo-pair skill). The tool's output is
 * the scratch run's text (stdout/result + code mode's summary).
 */
export const RegisterLanguageModelToolsLive = Layer.scopedDiscard(
  Effect.gen(function* () {
    const code = yield* VsCode;
    const notebooks = yield* NotebookRuntime;
    const runPromise = Runtime.runPromise(yield* Effect.runtime());
    const decoder = new TextDecoder();

    /**
     * Resolve the explicit URI against open notebooks
     */
    const resolveNotebookId = Effect.fn(function* (
      input: typeof ExecuteCodeInput.Type,
    ) {
      const notebooks = yield* code.workspace.getNotebookDocuments();

      const first = EffectArray.findFirst(
        EffectArray.getSomes(
          notebooks.map((raw) => MarimoNotebookDocument.tryFrom(raw)),
        ),
        (notebook) => notebook.id === input.notebookUri,
      );

      return Option.map(first, (nb) => nb.id);
    });

    const result = (text: string) =>
      new code.LanguageModelToolResult([new code.LanguageModelTextPart(text)]);

    const executeCode = Effect.fn("lm.executeCode")(function* (
      unknownInput: unknown,
    ) {
      const input = yield* Schema.decodeUnknown(ExecuteCodeInput)(unknownInput);

      const notebookId = yield* resolveNotebookId(input);
      if (Option.isNone(notebookId)) {
        return result(
          `No open marimo notebook matches \`${input.notebookUri}\`. ` +
            `Pass the URI of a marimo notebook open in VS Code's notebook editor. ` +
            `Do not fall back to editing the \`.py\` file with Edit/Write/NotebookEdit — ` +
            `that bypasses the live kernel.`,
        );
      }

      const ops = yield* notebooks
        .forNotebook(notebookId.value)
        .executeScratchpad(input.code)
        .pipe(Stream.runCollect);

      // Mirror marimo's SSE `/execute` endpoint (`ScratchCellListener.stream`):
      // surface the scratch cell's own output PLUS console (stdout/stderr) from
      // any cells a code-mode cascade re-ran
      const [scratchOps, cascadeOps] = EffectArray.partition(
        Chunk.toReadonlyArray(ops),
        (op) => extractCellIdFromCellMessage(op) === SCRATCH_CELL_ID,
      );

      // The scratch cell's own output (its console + rendered value).
      const scratchText = scratchCellNotificationsToVsCodeOutput(
        scratchOps,
        code,
      ).pipe(
        Option.map((cellOutput) =>
          cellOutput.items.map((item) => decoder.decode(item.data)).join(""),
        ),
        Option.getOrElse(() => ""),
      );

      // Console from cascade cells, concatenated in arrival order
      const cascadeText = cascadeOps.map(consoleText).join("");

      const text = scratchText + cascadeText;

      return result(text.trim() === "" ? "(no output)" : text);
    });

    yield* code.lm.registerTool<unknown>(EXECUTE_CODE_TOOL, {
      prepareInvocation(options) {
        const input = Schema.decodeUnknownSync(ExecuteCodeInput)(options.input);
        return {
          invocationMessage: "Running code in the marimo kernel…",
          // Side-effecting (arbitrary code in the user's kernel) — confirm.
          confirmationMessages: {
            title: "Run code in the marimo kernel?",
            message: new code.MarkdownString(
              `\`\`\`python\n${input.code}\n\`\`\``,
            ),
          },
        };
      },
      invoke(options, token) {
        return runPromise(
          executeCode(options.input).pipe(
            Effect.catchTags({
              // No live kernel — the notebook is open but nothing can run. Tell
              // the agent to have the user start one rather than file-edit.
              NoActiveKernelError: () =>
                Effect.succeed(
                  result(
                    `This marimo notebook is open but has no active kernel, so code can't run. ` +
                      `Ask the user to select a kernel and run a cell to start it, then try again. ` +
                      `Do not edit the \`.py\` file directly as a fallback — that bypasses the live kernel.`,
                  ),
                ),
              // Sandbox kernel needs a script file on disk to resolve its venv.
              UnsavedNotebookError: () =>
                Effect.succeed(
                  result(
                    `This notebook uses a sandbox kernel, which requires the notebook to be saved to a file before it can run. ` +
                      `Ask the user to save it, then try again.`,
                  ),
                ),
            }),
          ),
          { signal: signalFromToken(token) },
        );
      },
    });
  }),
);
