import * as NodeChildProcess from "node:child_process";
import * as NodeFs from "node:fs";
import * as NodePath from "node:path";
import * as NodeProcess from "node:process";

import {
  Cause,
  Data,
  Effect,
  Option,
  PubSub,
  Redacted,
  Runtime,
  Stream,
} from "effect";
import * as lsp from "vscode-languageclient/node";

import { Config } from "../config/Config.ts";
import { NOTEBOOK_TYPE } from "../constants.ts";
import { acquireDisposable } from "../lib/acquireDisposable.ts";
import { showErrorAndPromptLogs } from "../lib/showErrorAndPromptLogs.ts";
import { tokenFromSignal } from "../lib/tokenFromSignal.ts";
import { VsCode } from "../platform/VsCode.ts";
import { Uv } from "../python/Uv.ts";
import * as Api from "../schemas/Models.gen.ts";
import { Telemetry } from "../telemetry/Telemetry.ts";
import type {
  MarimoApiCall,
  MarimoOperation,
  MarimoSessionsChanged,
} from "../types.ts";

const MAX_STDERR_LINES = 200;

export type MarimoLspMode = "wasm" | "uv" | "configured";

const MarimoLspExecutable = Data.taggedEnum<MarimoLspExecutable>();
type MarimoLspExecutable = Data.TaggedEnum<{
  Configured: { readonly exec: lsp.Executable };
  Wasm: { readonly exec: lsp.Executable };
  Uv: { readonly exec: lsp.Executable };
}>;

export interface LspProcessExit {
  readonly code: number | null;
  readonly signal: NodeJS.Signals | null;
}

export class MarimoClientStartError extends Data.TaggedError(
  "MarimoClientStartError",
)<{
  exec: lsp.Executable;
  cause: unknown;
  /** Tail of the marimo-lsp stderr captured up to the failure, if any. */
  stderr?: string;
  /** Exit code/signal if marimo-lsp exited before or during startup. */
  exit?: LspProcessExit;
  mode: MarimoLspMode;
}> {}

export class MarimoCommandError extends Data.TaggedError("MarimoCommandError")<{
  readonly command: Redacted.Redacted<{
    readonly command: "marimo.api";
    readonly params: MarimoApiCall;
  }>;
  readonly cause: unknown;
  readonly mode: MarimoLspMode;
}> {}

/**
 * The small transport seam used to construct MarimoClient.
 *
 * Production uses vscode-languageclient. Tests can provide an in-memory
 * transport without reproducing MarimoClient's named methods.
 */
interface MarimoTransport<Error = never> {
  readonly execute: (request: MarimoApiCall) => Effect.Effect<unknown, Error>;
  readonly operations: () => Stream.Stream<MarimoOperation>;
  readonly sessionChanges?: () => Stream.Stream<MarimoSessionsChanged>;
}

export function makeMarimoCommands<Error>(transport: MarimoTransport<Error>) {
  return {
    operations: transport.operations,
    sessionChanges: transport.sessionChanges ?? (() => Stream.never),
    ...Api.makeApiClient(transport.execute),
  };
}

export const makeMarimoOperationStream = Effect.fn(function* (
  register: (handler: (message: MarimoOperation) => void) => {
    readonly dispose: () => void;
  },
) {
  const operationPubSub = yield* PubSub.unbounded<MarimoOperation>();
  const runSync = Runtime.runSync(yield* Effect.runtime());
  yield* Effect.addFinalizer(() => PubSub.shutdown(operationPubSub));

  // vscode-languageclient stores one notification handler per method, so
  // register once and fan out to every operations() consumer.
  yield* acquireDisposable(() =>
    register((message) => {
      runSync(PubSub.publish(operationPubSub, message));
    }),
  );

  return () => Stream.fromPubSub(operationPubSub);
});

/**
 * Communication with marimo-lsp.
 *
 * This module owns the marimo-lsp process, LSP transport, named commands,
 * and operation stream.
 */
export class MarimoClient extends Effect.Service<MarimoClient>()(
  "MarimoClient",
  {
    dependencies: [Config.Default, Uv.Default, Telemetry.Default],
    scoped: Effect.gen(function* () {
      const code = yield* VsCode;
      const config = yield* Config;
      const telemetry = yield* Telemetry;

      const configuredExec = yield* config.lsp.executable;
      const useWasm = yield* config.lsp.wasm;
      const uvBinary =
        Option.isNone(configuredExec) && !useWasm
          ? (yield* (yield* Uv).bin).executable
          : undefined;
      const selection = yield* selectMarimoLspExecutable({
        configuredExec,
        useWasm,
        uvBinary,
      });
      const { exec } = selection;
      const mode = marimoLspMode(selection);
      yield* telemetry.lspModeSelected(mode);

      yield* Effect.logInfo("Starting marimo-lsp").pipe(
        Effect.annotateLogs(
          MarimoLspExecutable.$match(selection, {
            Configured: ({ exec }) => ({
              mode: "configured",
              command: exec.command,
              args: (exec.args ?? []).join(" "),
            }),
            Wasm: () => ({ mode: "wasm" }),
            Uv: () => ({ mode: "uv" }),
          }),
        ),
      );

      const outputChannel =
        yield* code.window.createLogOutputChannel("marimo-lsp");

      interface SpawnState {
        readonly stderrTail: string[];
        pending: string;
        exit?: LspProcessExit;
      }

      // vscode-languageclient may start a replacement process before its
      // start promise rejects. Keep both references so startup errors report
      // stderr from the process that actually exited.
      let currentSpawn: SpawnState | undefined;
      let lastExitedSpawn: SpawnState | undefined;

      const serverOptions: lsp.ServerOptions = () =>
        new Promise<NodeChildProcess.ChildProcess>((resolve, reject) => {
          const spawn: SpawnState = { stderrTail: [], pending: "" };
          currentSpawn = spawn;

          const child = NodeChildProcess.spawn(
            exec.command,
            exec.args ?? [],
            exec.options,
          );
          child.once("error", reject);
          if (child.pid === undefined) return;

          child.stderr?.setEncoding("utf8");
          child.stderr?.on("data", (chunk: string) => {
            spawn.pending += chunk;
            let newline: number;
            while ((newline = spawn.pending.indexOf("\n")) !== -1) {
              let line = spawn.pending.slice(0, newline);
              spawn.pending = spawn.pending.slice(newline + 1);
              if (line.endsWith("\r")) line = line.slice(0, -1);
              if (line.length === 0) continue;
              spawn.stderrTail.push(line);
              if (spawn.stderrTail.length > MAX_STDERR_LINES) {
                spawn.stderrTail.shift();
              }
            }
          });
          child.once("exit", (code, signal) => {
            if (spawn.pending.length > 0) {
              spawn.stderrTail.push(spawn.pending);
              spawn.pending = "";
            }
            spawn.exit = { code, signal };
            lastExitedSpawn = spawn;
          });
          resolve(child);
        });

      const client = new lsp.LanguageClient(
        "marimo-lsp",
        "Marimo Language Server",
        serverOptions,
        {
          outputChannel,
          revealOutputChannelOn: lsp.RevealOutputChannelOn.Never,
          documentSelector: [
            { notebook: NOTEBOOK_TYPE, language: "sql" },
            { notebook: NOTEBOOK_TYPE, language: "python" },
            { notebook: NOTEBOOK_TYPE, language: "mo-python" },
            { notebook: NOTEBOOK_TYPE, language: "markdown" },
          ],
        },
      );

      const stopClient = Effect.fn(function* () {
        // LanguageClient.needsStop() also returns true while Starting, but
        // stop() rejects unless the client has an active Running connection.
        if (!client.isRunning()) return;
        yield* Effect.tryPromise(() => client.stop()).pipe(
          Effect.timeout("5 seconds"),
          Effect.ignore,
        );
        yield* Effect.logDebug("marimo-lsp client stopped");
      });

      const startClient = () =>
        Effect.gen(function* () {
          // start() is single-flight while the client is Starting, so there is
          // no need to stop an in-progress start before awaiting it.
          yield* Effect.tryPromise({
            try: () => client.start(),
            catch: (cause) => {
              const source =
                currentSpawn?.exit !== undefined
                  ? currentSpawn
                  : (lastExitedSpawn ?? currentSpawn);
              return new MarimoClientStartError({
                exec,
                mode,
                cause,
                stderr:
                  source && source.stderrTail.length > 0
                    ? source.stderrTail.join("\n")
                    : undefined,
                exit: source?.exit,
              });
            },
          });
          yield* telemetry.lspStarted(mode);
          yield* Effect.logInfo("marimo-lsp client started").pipe(
            Effect.annotateLogs({ "lsp.mode": mode }),
          );
        }).pipe(Effect.withSpan("lsp.start"));

      yield* Effect.addFinalizer(() => disposeLanguageClient(client));

      const operations = yield* makeMarimoOperationStream((handler) =>
        client.onNotification("marimo/operation", (message) => {
          handler(message);
        }),
      );

      const restart = () =>
        code.window.withProgress(
          {
            location: code.ProgressLocation.Notification,
            title: "Restarting marimo-lsp",
            cancellable: true,
          },
          Effect.fn(function* (progress) {
            if (client.isRunning()) {
              progress.report({ message: "Stopping..." });
              yield* stopClient();
            }
            progress.report({ message: "Starting..." });
            yield* startClient().pipe(
              Effect.tap(() =>
                Effect.sync(() => progress.report({ message: "Done." })),
              ),
              Effect.catchTag(
                "MarimoClientStartError",
                Effect.fn(function* (error) {
                  const message = "Failed to restart marimo-lsp.";
                  yield* Effect.logError(message).pipe(
                    Effect.annotateLogs({
                      cause: Cause.fail(error),
                      "lsp.mode": mode,
                    }),
                  );
                  yield* showErrorAndPromptLogs(message, {
                    channel: outputChannel,
                  });
                }),
              ),
            );
          }),
        );

      const transport: MarimoTransport<
        MarimoClientStartError | MarimoCommandError
      > = {
        execute: Effect.fn(function* (request) {
          if (!client.isRunning()) {
            yield* startClient();
          }
          const command = {
            command: "marimo.api",
            params: request,
          } as const;
          return yield* Effect.tryPromise({
            try: (signal) =>
              client.sendRequest<unknown>(
                "workspace/executeCommand",
                { command: command.command, arguments: [command.params] },
                tokenFromSignal(signal),
              ),
            catch: (cause) =>
              new MarimoCommandError({
                command: Redacted.make(command),
                cause,
                mode,
              }),
          }).pipe(
            Effect.withSpan("lsp.executeCommand", {
              attributes: {
                command: command.command,
                method: request.method,
              },
            }),
          );
        }),
        operations,
        sessionChanges: () =>
          Stream.asyncPush<MarimoSessionsChanged>((emit) =>
            acquireDisposable(() =>
              client.onNotification("marimo/sessionsChanged", (message) => {
                emit.single(message);
              }),
            ),
          ),
      };

      return {
        mode,
        channel: {
          name: outputChannel.name,
          show: outputChannel.show.bind(outputChannel),
        },
        restart,
        ...makeMarimoCommands(transport),
      };
    }),
  },
) {}

function marimoLspMode(selection: MarimoLspExecutable): MarimoLspMode {
  return MarimoLspExecutable.$match(selection, {
    Configured: () => "configured" as const,
    Wasm: () => "wasm" as const,
    Uv: () => "uv" as const,
  });
}

/**
 * Dispose the language client without allowing dependency cleanup failures to
 * fail the owning Effect scope. NodeLanguageClient still runs its process-kill
 * cleanup in a `finally`, even when graceful protocol shutdown rejects.
 */
export function disposeLanguageClient(
  client: Pick<lsp.LanguageClient, "dispose">,
): Effect.Effect<void> {
  return Effect.tryPromise(() => client.dispose()).pipe(
    Effect.timeout("5 seconds"),
    Effect.ignore,
  );
}

export const findMarimoLspExecutable = Effect.fn("findMarimoLspExecutable")(
  function* (uvBinary: string, searchDirectory = __dirname) {
    const sdistDir = NodeFs.readdirSync(searchDirectory).find((file) =>
      file.startsWith("marimo_lsp-"),
    );

    if (sdistDir) {
      const sdist = NodePath.join(searchDirectory, sdistDir);
      yield* Effect.logDebug("Using bundled marimo-lsp").pipe(
        Effect.annotateLogs({ sdist }),
      );
      return {
        command: uvBinary,
        args: [
          "tool",
          "run",
          "--python",
          ">=3.13,<3.15",
          "--from",
          sdist,
          "marimo-lsp",
        ],
      };
    }

    yield* Effect.logWarning("No bundled wheel found, using development mode");

    return {
      command: uvBinary,
      args: ["run", "--directory", searchDirectory, "marimo-lsp"],
    };
  },
);

export const selectMarimoLspExecutable = Effect.fn("selectMarimoLspExecutable")(
  function* ({
    configuredExec,
    useWasm,
    uvBinary,
    searchDirectory = __dirname,
  }: {
    readonly configuredExec: Option.Option<lsp.Executable>;
    readonly useWasm: boolean;
    readonly uvBinary?: string;
    readonly searchDirectory?: string;
  }) {
    if (Option.isSome(configuredExec)) {
      return MarimoLspExecutable.Configured({ exec: configuredExec.value });
    }
    if (useWasm) {
      return MarimoLspExecutable.Wasm({
        exec: findWasmMarimoLspExecutable(searchDirectory),
      });
    }
    if (uvBinary === undefined) {
      throw new Error(
        "uv is required when the WASM language server is disabled",
      );
    }
    return MarimoLspExecutable.Uv({
      exec: yield* findMarimoLspExecutable(uvBinary, searchDirectory),
    });
  },
);

export function findWasmMarimoLspExecutable(
  searchDirectory = __dirname,
): lsp.Executable {
  return {
    command: NodeProcess.execPath,
    args: [NodePath.join(searchDirectory, "wasmServer.js")],
    options: {
      env: {
        ...NodeProcess.env,
        // VS Code runs extensions in Electron. This makes the dedicated
        // language-server child use Electron's bundled Node runtime.
        ELECTRON_RUN_AS_NODE: "1",
      },
    },
  };
}
