import * as semver from "@std/semver";
import type * as py from "@vscode/python-extension";
import {
  Brand,
  Cause,
  Effect,
  Option,
  Redacted,
  Runtime,
  Stream,
} from "effect";
import type * as vscode from "vscode";

import { unreachable } from "../assert.ts";
import { Config } from "../config/Config.ts";
import { acquireDisposable } from "../lib/acquireDisposable.ts";
import { extractExecuteCodeRequest } from "../lib/extractExecuteCodeRequest.ts";
import { extractPythonError } from "../lib/extractPythonError.ts";
import { formatControllerLabel } from "../lib/formatControllerLabel.ts";
import { installPackages } from "../lib/installPackages.ts";
import { isProblematicFilename } from "../lib/validateNotebookFilename.ts";
import { NotebookSerializer } from "../notebook/NotebookSerializer.ts";
import { Constants } from "../platform/Constants.ts";
import { VsCode } from "../platform/VsCode.ts";
import { EnvironmentValidator } from "../python/EnvironmentValidator.ts";
import { findVenvPath } from "../python/findVenvPath.ts";
import { Uv } from "../python/Uv.ts";
import {
  type MarimoNotebookCell,
  MarimoNotebookDocument,
} from "../schemas/MarimoNotebookDocument.ts";
import { NotebookRuntime } from "./NotebookRuntime.ts";

const NotebookControllerId = Brand.nominal<NotebookControllerId>();
export type NotebookControllerId = Brand.Branded<string, "ControllerId">;

export const createPythonController = Effect.fn("createPythonController")(
  function* (options: {
    id: NotebookControllerId;
    label: string;
    env: py.Environment;
  }) {
    const uv = yield* Uv;
    const code = yield* VsCode;
    const config = yield* Config;
    const notebooks = yield* NotebookRuntime;
    const validator = yield* EnvironmentValidator;
    const serializer = yield* NotebookSerializer;
    const { LanguageId } = yield* Constants;
    const runPromise = Runtime.runPromise(yield* Effect.runtime());

    yield* Effect.annotateCurrentSpan("controllerId", options.id);
    const controller = yield* code.notebooks.createNotebookController(
      options.id,
      serializer.notebookType,
      options.label,
    );

    // Add metadata
    controller.supportedLanguages = [LanguageId.Python, LanguageId.Sql];
    controller.description = options.env.path;

    // Set up execution handler
    controller.executeHandler = (rawCells, rawNotebook, controller) =>
      runPromise(
        Effect.gen(function* () {
          const request = extractExecuteCodeRequest(rawCells, LanguageId);
          if (Option.isNone(request)) {
            yield* Effect.logWarning("Empty execution request");
            return;
          }

          const notebook = MarimoNotebookDocument.from(rawNotebook);

          const validation = isProblematicFilename(rawNotebook.uri);
          if (validation.problematic) {
            yield* code.window.showErrorMessage(validation.message, {
              modal: true,
            });
            return;
          }

          const validEnv = yield* validator.validate(options.env);

          yield* notebooks
            .forNotebook(notebook.id)
            .executeCells(request.value, validEnv.executable);
        }).pipe(
          Effect.withSpan("PythonController.execute", {
            attributes: {
              controllerId: controller.id,
              cellCount: rawCells.length,
              notebook: rawNotebook.uri.toString(),
            },
          }),
          // Known exceptions
          Effect.catchTags({
            NotebookFileRootError: Effect.fn(function* (error) {
              yield* Effect.logError(error.message).pipe(
                Effect.annotateLogs({ configuredValue: error.configuredValue }),
              );
              yield* code.window.showErrorMessage(error.message, {
                modal: true,
              });
            }),
            NoActiveKernelError: Effect.fn(function* () {
              yield* code.window.showErrorMessage(
                "The notebook was closed before its kernel could start.",
                { modal: true },
              );
            }),
            MarimoCommandError: Effect.fn(function* (error) {
              yield* Effect.logError("Failed to execute command").pipe(
                Effect.annotateLogs({
                  cause: Cause.fail(error),
                  command: Redacted.value(error.command).command,
                }),
              );
              const detail = extractPythonError(error.cause);
              yield* code.window.showErrorMessage(
                Option.isSome(detail)
                  ? `Failed to execute marimo command:\n\n${detail.value}`
                  : "Failed to execute marimo command. Please check the logs for details.",
                { modal: true },
              );
            }),
            EnvironmentInspectionError: Effect.fn(function* (error) {
              yield* Effect.logError("Python venv check failed").pipe(
                Effect.annotateLogs({
                  cause: Cause.fail(error),
                  pythonPath: error.env.path,
                  stdout: error.stdout,
                  stderr: error.stderr,
                }),
              );

              if (error.cause?._tag === "InvalidExecutableError") {
                yield* code.window.showErrorMessage(
                  `Python executable does not exist for env: ${error.env.path}.`,
                  { modal: true },
                );
              } else {
                const stderrSnippet = error.stderr
                  ? `\n\nstderr:\n${truncate(error.stderr.trim(), 500)}`
                  : "";
                yield* code.window.showErrorMessage(
                  `Failed to check dependencies in ${formatControllerLabel(code, options.env)}.\n\n` +
                    `Python path: ${error.env.path}` +
                    stderrSnippet,
                  { modal: true },
                );
              }
            }),
            EnvironmentRequirementError: Effect.fn(function* (error) {
              yield* Effect.logWarning("Environment requirements not met").pipe(
                Effect.annotateLogs({
                  pythonPath: error.env.path,
                  diagnostics: error.diagnostics,
                }),
              );
              const messages = error.diagnostics.map((d) => {
                switch (d.kind) {
                  case "missing":
                    return `• ${d.package}: not installed`;
                  case "outdated":
                    return `• ${d.package}: v${semver.format(d.currentVersion)} (requires >=v${semver.format(d.requiredVersion)})`;
                  case "unknown":
                    return `• ${d.package}: unable to detect`;
                  default:
                    return unreachable(d);
                }
              });

              // Only prompt to install if uv is enabled and we have a venv
              // Non-venv environments (pixi, conda, bazel, global) don't have pyvenv.cfg
              // so uv can't install packages there
              const venv = findVenvPath(options.env.path);
              const uvEnabled = yield* config.uv.enabled;
              const canInstallWithUv = uvEnabled && Option.isSome(venv);

              if (canInstallWithUv) {
                const msg =
                  `${formatControllerLabel(code, options.env)} cannot run the marimo kernel:\n\n` +
                  messages.join("\n") +
                  `\n\nPackages are missing or outdated.\n\nInstall with uv?`;

                const choice = yield* code.window.showErrorMessage(msg, {
                  modal: true,
                  items: ["Yes"],
                });
                if (!choice) {
                  return;
                }
                const packages = error.diagnostics.map((d) =>
                  d.kind === "outdated"
                    ? `${d.package}>=${semver.format(d.requiredVersion)}`
                    : d.package,
                );
                yield* installPackages(packages, {
                  venvPath: venv.value,
                }).pipe(
                  Effect.provideService(VsCode, code),
                  Effect.provideService(Uv, uv),
                );
              } else {
                const msg =
                  `${formatControllerLabel(code, options.env)} cannot run the marimo kernel:\n\n` +
                  messages.join("\n") +
                  `\n\nPlease install or update the missing packages.`;

                yield* code.window.showErrorMessage(msg, {
                  modal: true,
                });
              }
            }),
          }),
        ),
      );

    // Set up interrupt handler
    controller.interruptHandler = (rawNotebook) =>
      runPromise(
        Effect.gen(function* () {
          const notebook = MarimoNotebookDocument.from(rawNotebook);
          yield* notebooks.forNotebook(notebook.id).interrupt();
        }).pipe(
          Effect.withSpan("PythonController.interrupt", {
            attributes: {
              controllerId: controller.id,
              notebook: rawNotebook.uri.toString(),
            },
          }),
          Effect.catchAllCause((cause) =>
            Effect.gen(function* () {
              yield* Effect.logError("Failed to interrupt execution").pipe(
                Effect.annotateLogs({ cause }),
              );
              yield* code.window.showErrorMessage(
                "Failed to interrupt execution. Please check the logs for details.",
              );
            }),
          ),
        ),
      );

    return new PythonController(controller, options.env.path);
  },
);

export class PythonController {
  readonly _tag = "PythonController";
  #inner: Omit<vscode.NotebookController, "dispose">;
  /** The python interpreter this controller's environment runs on. */
  executable: string;
  constructor(
    inner: Omit<vscode.NotebookController, "dispose">,
    executable: string,
  ) {
    this.#inner = inner;
    this.executable = executable;
  }
  static getId(env: py.Environment) {
    return NotebookControllerId(`marimo-${env.path}`);
  }
  get id(): NotebookControllerId {
    // SAFETY: `this.#inner` is constructed with an id produced by the static
    // getId() helper below, which runs the string through the
    // NotebookControllerId brand constructor.
    // oxlint-disable-next-line typescript/no-unsafe-type-assertion
    return this.#inner.id as NotebookControllerId;
  }
  mutateDescription(description: string) {
    return Effect.sync(() => {
      this.#inner.description = description;
      return this;
    });
  }
  createNotebookCellExecution(cell: MarimoNotebookCell) {
    return this.#inner.createNotebookCellExecution(cell.rawNotebookCell);
  }
  resolveExecutable(_notebook: MarimoNotebookDocument) {
    return Effect.succeed(this.executable);
  }
  selectedNotebookChanges() {
    return Stream.asyncPush<{
      notebook: vscode.NotebookDocument;
      selected: boolean;
    }>((emit) =>
      acquireDisposable(() =>
        this.#inner.onDidChangeSelectedNotebooks((e) => emit.single(e)),
      ),
    );
  }
  updateNotebookAffinity(
    notebook: vscode.NotebookDocument,
    affinity: vscode.NotebookControllerAffinity,
  ) {
    return Effect.sync(() => {
      this.#inner.updateNotebookAffinity(notebook, affinity);
    });
  }
}

function truncate(text: string, maxLength: number): string {
  if (text.length <= maxLength) {
    return text;
  }
  return `${text.slice(0, maxLength)}…`;
}
