import { join, resolve, sep } from "node:path";
import { pathToFileURL } from "node:url";
import type { BridgeConnectionConfig } from "../../bridge/protocol.ts";
import {
  RESERVED_AGENT_TOOL,
  RESERVED_APPROVAL_TOOL,
  RESERVED_LOAD_TOOL,
  RESERVED_MCP_TOOL,
  RESERVED_MEMORY_TOOL,
  RESERVED_OUTPUT_TOOL,
  RESERVED_SCHEMA_TOOL,
  RESERVED_STORE_TOOL,
  TIMEOUT_PAUSE_OP,
  TIMEOUT_RESUME_OP,
} from "../../bridge/reserved.ts";
import type { JavaScriptKernelOptions as BaseJavaScriptKernelOptions } from "./kernel-contract.ts";
import { rewriteImports } from "./rewrite-imports.ts";

const PREPARED_CELL_PREFIX = "/*pi-codemode:prepared-cell*/";

export interface LocalModuleLoaderOptions {
  readonly artifactsDir?: string;
  readonly bridgeTimeoutSeconds?: number;
  readonly cwd: string;
  readonly localRoots?: Readonly<Record<string, string>>;
  readonly toolNames?: readonly string[];
  readonly tools?: readonly {
    readonly description?: string;
    readonly name: string;
  }[];
}

export type JavaScriptKernelOptions = BaseJavaScriptKernelOptions &
  LocalModuleLoaderOptions;

const DEFAULT_BRIDGE_TIMEOUT_SECONDS = 120;

export function localBridgeConnection(
  options: LocalModuleLoaderOptions
): BridgeConnectionConfig {
  return {
    port: 1,
    token: "local",
    bridgeTimeoutMs: Math.max(
      1,
      Math.round(
        (options.bridgeTimeoutSeconds ?? DEFAULT_BRIDGE_TIMEOUT_SECONDS) * 1000
      )
    ),
    ...(options.localRoots ? { localRoots: { ...options.localRoots } } : {}),
    ...(options.artifactsDir ? { artifactsDir: options.artifactsDir } : {}),
    ...(options.toolNames ? { toolNames: [...options.toolNames] } : {}),
    ...(options.tools
      ? { tools: options.tools.map((tool) => ({ ...tool })) }
      : {}),
  };
}

interface RuntimeModuleContext {
  readonly cwdUrl: string;
  readonly localRootUrls: Readonly<Record<string, string>>;
  readonly reservedAgentTool: string;
  readonly reservedApprovalTool: string;
  readonly reservedLoadTool: string;
  readonly reservedMcpTool: string;
  readonly reservedMemoryTool: string;
  readonly reservedOutputTool: string;
  readonly reservedSchemaTool: string;
  readonly reservedStoreTool: string;
  readonly timeoutPauseOp: string;
  readonly timeoutResumeOp: string;
}

function directoryUrl(directory: string): string {
  return pathToFileURL(`${resolve(directory)}${sep}`).href;
}

function runtimeContext(
  options: LocalModuleLoaderOptions
): RuntimeModuleContext {
  const roots: Record<string, string> = {};
  for (const [scheme, root] of Object.entries(options.localRoots ?? {})) {
    roots[scheme.toLowerCase()] = directoryUrl(root);
  }
  if (options.artifactsDir && roots.local === undefined) {
    roots.local = directoryUrl(join(options.artifactsDir, "local"));
  }
  return {
    cwdUrl: directoryUrl(options.cwd),
    localRootUrls: roots,
    reservedAgentTool: RESERVED_AGENT_TOOL,
    reservedApprovalTool: RESERVED_APPROVAL_TOOL,
    reservedLoadTool: RESERVED_LOAD_TOOL,
    reservedMcpTool: RESERVED_MCP_TOOL,
    reservedMemoryTool: RESERVED_MEMORY_TOOL,
    reservedOutputTool: RESERVED_OUTPUT_TOOL,
    reservedSchemaTool: RESERVED_SCHEMA_TOOL,
    reservedStoreTool: RESERVED_STORE_TOOL,
    timeoutPauseOp: TIMEOUT_PAUSE_OP,
    timeoutResumeOp: TIMEOUT_RESUME_OP,
  };
}

function loaderPrelude(context: RuntimeModuleContext): string {
  const serialized = JSON.stringify(context);
  return [
    `globalThis.__codemode_module_context__ = ${serialized};`,
    "globalThis.__codemode_reserved_agent_tool__ = globalThis.__codemode_module_context__.reservedAgentTool;",
    "globalThis.__codemode_reserved_output_tool__ = globalThis.__codemode_module_context__.reservedOutputTool;",
    "globalThis.__codemode_reserved_schema_tool__ = globalThis.__codemode_module_context__.reservedSchemaTool;",
    "globalThis.__codemode_reserved_store_tool__ = globalThis.__codemode_module_context__.reservedStoreTool;",
    "globalThis.__codemode_reserved_load_tool__ = globalThis.__codemode_module_context__.reservedLoadTool;",
    "globalThis.__codemode_reserved_mcp_tool__ = globalThis.__codemode_module_context__.reservedMcpTool;",
    "globalThis.__codemode_reserved_memory_tool__ = globalThis.__codemode_module_context__.reservedMemoryTool;",
    "globalThis.__codemode_reserved_approval_tool__ = globalThis.__codemode_module_context__.reservedApprovalTool;",
    "globalThis.__codemode_timeout_pause_op__ = globalThis.__codemode_module_context__.timeoutPauseOp;",
    "globalThis.__codemode_timeout_resume_op__ = globalThis.__codemode_module_context__.timeoutResumeOp;",
    "globalThis.__codemode_import__ = async (source, options) => {",
    "  const context = globalThis.__codemode_module_context__;",
    "  const specifier = String(source);",
    "  const match = /^([a-z][a-z0-9+.-]*):\\/\\/(.*)$/i.exec(specifier);",
    "  let target = specifier;",
    "  if (match) {",
    "    const scheme = match[1].toLowerCase();",
    "    const root = context.localRootUrls[scheme];",
    "    if (!root) throw new Error('Unsupported module protocol: ' + specifier);",
    "    let relative;",
    "    try { relative = decodeURIComponent(match[2].replaceAll('\\\\', '/')); }",
    "    catch { throw new Error('Invalid module URL encoding: ' + specifier); }",
    "    if (relative.startsWith('/') || relative.split('/').includes('..')) {",
    "      throw new Error('Module path escapes ' + scheme + ':// root: ' + specifier);",
    "    }",
    "    target = new URL(relative, root).href;",
    "    const urlModule = await import('node:url');",
    "    const fsModule = await import('node:fs');",
    "    const pathModule = await import('node:path');",
    "    const rootPath = urlModule.fileURLToPath(root);",
    "    const filePath = urlModule.fileURLToPath(target);",
    "    const realpathPrefix = (candidate) => {",
    "      let ancestor = candidate;",
    "      const tail = [];",
    "      for (;;) {",
    "        try {",
    "          let resolved = fsModule.realpathSync(ancestor);",
    "          for (const part of tail) resolved = pathModule.join(resolved, part);",
    "          return resolved;",
    "        } catch (error) {",
    "          if (!error || error.code !== 'ENOENT') throw error;",
    "          const parent = pathModule.dirname(ancestor);",
    "          if (parent === ancestor) throw error;",
    "          tail.unshift(pathModule.basename(ancestor));",
    "          ancestor = parent;",
    "        }",
    "      }",
    "    };",
    "    const realRoot = realpathPrefix(rootPath);",
    "    const realTarget = realpathPrefix(filePath);",
    "    if (realTarget !== realRoot && !realTarget.startsWith(realRoot + pathModule.sep)) {",
    "      throw new Error('Module path escapes ' + scheme + ':// root: ' + specifier);",
    "    }",
    "  } else if (specifier.startsWith('./') || specifier.startsWith('../') || specifier === '.' || specifier === '..') {",
    "    target = new URL(specifier, context.cwdUrl).href;",
    "  } else if (specifier.startsWith('/') || /^[A-Za-z]:[\\\\/]/.test(specifier)) {",
    "    const urlModule = await import('node:url');",
    "    target = urlModule.pathToFileURL(specifier).href;",
    "  }",
    "  return options === undefined ? import(target) : import(target, options);",
    "};",
  ].join("\n");
}

export class LocalModuleLoader {
  readonly #prelude: string;

  constructor(options: LocalModuleLoaderOptions) {
    this.#prelude = loaderPrelude(runtimeContext(options));
  }

  prepareCell(code: string): string {
    return `${PREPARED_CELL_PREFIX}${JSON.stringify({ prelude: this.#prelude, code: rewriteImports(code) })}`;
  }
}