import { join } from "node:path";
import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import {
  type BridgeHttpCallRequest,
  type BridgeServerHandle,
  startBridgeServer,
} from "../bridge/http-server.ts";
import type { KernelToHostMessage } from "../bridge/protocol.ts";
import {
  isReservedToolName,
  runReservedTool,
} from "../bridges/reserved-dispatch.ts";
import type { EvalSchemaToolInfo } from "../bridges/schema-bridge.ts";
import { appendSchemaHint } from "../bridges/schema-hint.ts";
import type {
  CompletionRequest,
  CompletionResult,
} from "../completion/handler.ts";
import {
  type CodemodeSettings,
  defaultCodemodeSettings,
  resolveParallelPoolWidth,
} from "../config/settings.ts";
import type { InterpreterAvailability } from "../interpreters/detect.ts";
import { JavaScriptKernel } from "../kernels/js/context-manager.ts";
import { TypeScriptKernel } from "../kernels/ts/kernel.ts";
import { PythonKernel } from "../kernels/py/kernel.ts";
import { RubyKernel } from "../kernels/rb/kernel.ts";
import { sanitizeTerminalLabel } from "../render/sanitize.ts";
import {
  boundToolCallArgs,
  capCodePoints,
  toolCallResultPreview,
} from "../tool/call-capture.ts";
import { marshalToolResult, toolResultIsError } from "../tool/image.ts";
import { SessionStore } from "../tool/session-store.ts";
import type {
  BridgeCallSession,
  EvalKernel,
  EvalKernelManager,
  EvalLanguage,
  EvalToolCallSummary,
  ExecuteTool,
} from "../tool/types.ts";

export interface CodemodeSessionManager extends EvalKernelManager {
  bridgeEndpoint?: () => BridgeEndpoint;
  complete: (
    request: CompletionRequest,
    ctx: ExtensionContext
  ) => Promise<CompletionResult>;
  dispose: () => Promise<void>;
  /**
   * Register (or with `undefined`, unregister) the per-active-cell bridge
   * session for a language: the HTTP /call route records tool-call enrichment
   * and forwards agent() status events through it, mirroring the JS kernel
   * path. One active cell per language is the invariant the detached cell
   * manager enforces, so a single slot per language suffices.
   */
  setBridgeCellSession?: (
    language: EvalLanguage,
    session: BridgeCallSession | undefined
  ) => void;
  setContext?: (ctx: ExtensionContext) => void;
}

export interface BridgeEndpoint {
  readonly port: number;
  readonly token: string;
}

export interface EvalExecutionTracker {
  assertEvalExecutionAllowed: () => void;
  trackEvalExecution: <Result>(
    execution: Promise<Result>,
    controller: AbortController
  ) => Promise<Result>;
}

export interface CreateCodemodeSessionManagerOptions {
  /** Session-adjacent directory used for persisted eval artifacts. */
  readonly artifactsDir?: string;
  readonly availability: InterpreterAvailability;
  readonly complete: (
    request: CompletionRequest,
    ctx: ExtensionContext
  ) => Promise<CompletionResult>;
  readonly cwd: string;
  readonly executeTool: ExecuteTool;
  readonly listTools?: () => readonly EvalSchemaToolInfo[];
  /** Session-scoped roots exposed to kernel helpers such as local://. */
  readonly localRoots?: Readonly<Record<string, string>>;
  readonly sessionId: string;
  /** Session value store backing the store()/load() helpers; defaults to a fresh store when absent. */
  readonly sessionStore?: SessionStore;
  readonly settings: CodemodeSettings;
}

export async function createCodemodeSessionManager(
  options: CreateCodemodeSessionManagerOptions
): Promise<CodemodeSessionManager> {
  const manager = new DefaultCodemodeSessionManager(options);
  await manager.start();
  return manager;
}

export class CodemodeSessionDisposedError extends Error {
  readonly name = "CodemodeSessionDisposedError";

  constructor() {
    super("codemode session manager is disposed");
  }
}

class CodemodeContextUnavailableError extends Error {
  readonly name = "CodemodeContextUnavailableError";

  constructor() {
    super("codemode completion context is unavailable");
  }
}

class DefaultCodemodeSessionManager implements CodemodeSessionManager {
  readonly #options: CreateCodemodeSessionManagerOptions;
  readonly #sessionStore: SessionStore;
  #bridge: BridgeServerHandle | undefined;
  readonly #kernels = new Map<EvalLanguage, EvalKernel>();
  readonly #kernelCreations = new Map<EvalLanguage, Promise<EvalKernel>>();
  readonly #onMessageRefs = new Map<
    EvalLanguage,
    (message: KernelToHostMessage) => void
  >();
  readonly #bridgeCellSessions = new Map<EvalLanguage, BridgeCallSession>();
  #context: ExtensionContext | undefined;
  // Generation counter for the session-replacement safety mechanism: must
  // mirror SessionManagerProxy.#generation (both are incremented on every
  // dispose/replacement). They guard different layers, so they are not merged.
  #generation = 0;
  #disposePromise: Promise<void> | undefined;

  constructor(options: CreateCodemodeSessionManagerOptions) {
    this.#options = options;
    this.#sessionStore = options.sessionStore ?? new SessionStore();
  }

  async start(): Promise<void> {
    this.#bridge = await startBridgeServer({
      onCall: async (request) => await this.#call(request),
      onCompletion: async (request) =>
        this.#options.complete(
          { prompt: request.prompt, opts: request.opts },
          this.#contextFor(request.signal)
        ),
    });
  }

  // Subprocess kernels (py/rb) reach the host only through this route, so reserved
  // helper names must dispatch exactly as the in-process JS path does in tool/cell-handler.ts.
  // Forwarding them to executeTool made agent() fail with "Unknown tool __agent__".
  //
  // Enrichment parity: the route records tool-call summaries (args, duration,
  // result preview; capped by the shared tracker in tool/call-capture.ts) and
  // agent() progress status events into the language's active-cell session
  // registered by the eval tool — the same pipeline the JS kernel path gets
  // from CellHandler. The HTTP reply shape stays { ok, value }-compatible with
  // what the preludes consume (they read value/error only); the metadata is
  // summary-only and never enters the model-visible reply value.
  async #call(request: BridgeHttpCallRequest): Promise<unknown> {
    const session =
      request.language === undefined
        ? undefined
        : this.#bridgeCellSessions.get(request.language);
    const startedAt = Date.now();
    const capturedArgs = boundToolCallArgs(request.args);
    const record = (
      ok: boolean,
      error: string | undefined,
      resultPreview: string | undefined
    ): void => {
      if (session === undefined) {
        return;
      }
      const summary: EvalToolCallSummary = {
        name: request.toolName,
        ok,
        ...(error === undefined ? {} : { error }),
      };
      // Reserved calls are recorded unenriched in the JS path too (no callId/
      // args/duration in CellHandler); mirror that for identical summaries.
      if (!isReservedToolName(request.toolName)) {
        Object.assign(summary, {
          callId: request.callId,
          args: capturedArgs.args,
          durationMs: Date.now() - startedAt,
          ...(capturedArgs.truncated ? { argsTruncated: true } : {}),
          ...(resultPreview === undefined ? {} : { resultPreview }),
        });
      }
      session.trackToolCall(summary);
    };
    try {
      if (!isReservedToolName(request.toolName)) {
        const result = await this.#options.executeTool(
          request.toolName,
          request.args,
          {
            signal: request.signal,
          }
        );
        if (toolResultIsError(result)) {
          let errorText: string | undefined;
          for (const part of result.content) {
            if (part.type !== "text") {
              continue;
            }
            errorText = capCodePoints(sanitizeTerminalLabel(part.text), 512);
            break;
          }
          record(false, errorText, undefined);
        } else {
          record(true, undefined, toolCallResultPreview(result));
        }
        return result;
      }
      const taskTools =
        this.#options.settings.taskTools ?? defaultCodemodeSettings.taskTools;
      const value = await runReservedTool(request.toolName, {
        callId: request.callId,
        args: request.args,
        executeTool: this.#options.executeTool,
        taskToolName: taskTools.task,
        taskOutputToolName: taskTools.output,
        listTools: this.#options.listTools,
        signal: request.signal,
        emitStatus:
          session === undefined
            ? () => {}
            : (event) => session.emitStatus(event),
        marshalToolResult,
        sessionStore: this.#sessionStore,
      });
      record(true, undefined, undefined);
      return value;
    } catch (error) {
      // Schema-hint parity with CellHandler.#deliverToolReply's catch path: the
      // error reply the prelude raises must carry the hint JS cells see.
      const text = appendSchemaHint(
        error instanceof Error ? error.message : String(error),
        request.toolName,
        this.#toolParameters(request.toolName)
      );
      record(false, text, undefined);
      if (error instanceof Error) {
        error.message = text;
        throw error;
      }
      throw new Error(text, { cause: error });
    }
  }

  setBridgeCellSession(
    language: EvalLanguage,
    session: BridgeCallSession | undefined
  ): void {
    if (session === undefined) {
      this.#bridgeCellSessions.delete(language);
    } else {
      this.#bridgeCellSessions.set(language, session);
    }
  }

  #toolParameters(toolName: string): unknown {
    return this.#options.listTools?.().find((tool) => tool.name === toolName)
      ?.parameters;
  }

  async getKernel(
    language: EvalLanguage,
    onMessage: (message: KernelToHostMessage) => void
  ): Promise<EvalKernel> {
    if (this.#disposePromise) {
      throw new CodemodeSessionDisposedError();
    }
    // Persistent kernels are reused across cells, but each cell needs its OWN
    // onMessage (bound to that cell's streaming state). Rebind on every call via
    // a stable dispatcher so the 2nd+ cell's text/display/log output is attributed
    // to the current cell, not the one that first created the kernel.
    this.#onMessageRefs.set(language, onMessage);
    const existing = this.#kernels.get(language);
    if (existing) {
      return existing;
    }
    const pending = this.#kernelCreations.get(language);
    if (pending) {
      return await pending;
    }
    const dispatch = (message: KernelToHostMessage): void =>
      this.#onMessageRefs.get(language)?.(message);
    const generation = this.#generation;
    const creation = this.#createAndStoreKernel(language, dispatch, generation);
    this.#kernelCreations.set(language, creation);
    try {
      return await creation;
    } finally {
      if (this.#kernelCreations.get(language) === creation) {
        this.#kernelCreations.delete(language);
      }
    }
  }

  async complete(
    request: CompletionRequest,
    ctx: ExtensionContext
  ): Promise<CompletionResult> {
    return await this.#options.complete(request, ctx);
  }

  bridgeEndpoint(): BridgeEndpoint {
    const bridge = this.#bridge;
    if (!bridge) {
      throw new Error("codemode bridge server is not running");
    }
    return { port: bridge.port, token: bridge.token };
  }

  setContext(ctx: ExtensionContext): void {
    this.#context = ctx;
  }

  dispose(): Promise<void> {
    if (this.#disposePromise) {
      return this.#disposePromise;
    }
    this.#generation += 1;
    this.#disposePromise = this.#disposeGeneration();
    return this.#disposePromise;
  }

  async #disposeGeneration(): Promise<void> {
    await Promise.allSettled(this.#kernelCreations.values());
    const kernels = [...this.#kernels.values()];
    const bridge = this.#bridge;
    this.#kernels.clear();
    this.#onMessageRefs.clear();
    this.#bridgeCellSessions.clear();
    this.#bridge = undefined;
    this.#context = undefined;
    const failures: unknown[] = [];
    for (const outcome of await Promise.allSettled(
      kernels.map((kernel) => kernel.close())
    )) {
      if (outcome.status === "rejected") {
        failures.push(outcome.reason);
      }
    }
    if (bridge) {
      const [outcome] = await Promise.allSettled([bridge.close()]);
      if (outcome?.status === "rejected") {
        failures.push(outcome.reason);
      }
    }
    if (failures.length > 0) {
      throw new AggregateError(
        failures,
        "Failed to dispose codemode session manager"
      );
    }
  }

  async #createAndStoreKernel(
    language: EvalLanguage,
    onMessage: (message: KernelToHostMessage) => void,
    generation: number
  ): Promise<EvalKernel> {
    const kernel = await this.#createKernel(language, onMessage);
    if (generation !== this.#generation) {
      await kernel.close();
      throw new CodemodeSessionDisposedError();
    }
    this.#kernels.set(language, kernel);
    return kernel;
  }

  async #createKernel(
    language: EvalLanguage,
    onMessage: (message: KernelToHostMessage) => void
  ): Promise<EvalKernel> {
    const bridge = this.#bridge;
    if (!bridge) {
      throw new Error("codemode bridge server is not running");
    }
    const parallelPoolWidth = resolveParallelPoolWidth(
      this.#options.settings.parallelPoolWidth
    );
    const bridgeTimeoutSeconds = this.#options.settings.bridgeTimeoutSeconds;
    const tools = this.#options.listTools?.().map((tool) => ({
      name: tool.name,
      ...(tool.description === undefined
        ? {}
        : { description: tool.description }),
    }));
    if (language === "js" || language === "ts") {
      const options = {
        sessionId: this.#options.sessionId,
        cwd: this.#options.cwd,
        parallelPoolWidth,
        bridgeTimeoutSeconds,
        hardenedCells: this.#options.settings.hardenedCells,
        jitless: this.#options.settings.jitless,
        ...(tools === undefined
          ? {}
          : { toolNames: tools.map((tool) => tool.name), tools }),
        onMessage,
      };
      return language === "js"
        ? new JavaScriptKernel(options)
        : new TypeScriptKernel(options);
    }
    const detected = this.#options.availability[language].detected;
    if (!detected.ok) {
      throw new Error(`No ${language} interpreter is available`);
    }
    const localRoots =
      this.#options.localRoots ??
      (this.#options.artifactsDir
        ? { local: join(this.#options.artifactsDir, "local") }
        : undefined);
    const connection = {
      port: bridge.port,
      token: bridge.token,
      parallelPoolWidth,
      bridgeTimeoutSeconds,
      ...(localRoots ? { localRoots: { ...localRoots } } : {}),
      ...(this.#options.artifactsDir
        ? { artifactsDir: this.#options.artifactsDir }
        : {}),
      ...(tools === undefined
        ? {}
        : { toolNames: tools.map((tool) => tool.name), tools }),
    };
    if (language === "py") {
      return await PythonKernel.start({
        interpreterPath: detected.path,
        sessionId: this.#options.sessionId,
        cwd: this.#options.cwd,
        connection,
        onMessage,
      });
    }
    if (language === "rb") {
      return RubyKernel.start({
        command: detected.path,
        sessionId: this.#options.sessionId,
        cwd: this.#options.cwd,
        connection,
        onMessage,
      });
    }
    throw new Error(`No kernel is available for ${language}`);
  }

  #contextFor(signal: AbortSignal): ExtensionContext {
    const ctx = this.#context;
    if (!ctx) {
      throw new CodemodeContextUnavailableError();
    }
    return {
      ...ctx,
      signal: ctx.signal ? AbortSignal.any([ctx.signal, signal]) : signal,
    };
  }
}