import type {
  RuntimeExecutionRequest,
  RuntimeExecutionResult,
  RuntimeWorkspaceFile,
} from '@learnitive/scientific-runtime';

const PYODIDE_VERSION = '314.0.4';
const PYODIDE_INDEX_URL = `https://cdn.jsdelivr.net/pyodide/v${PYODIDE_VERSION}/full/`;
const PYODIDE_MODULE_URL = `${PYODIDE_INDEX_URL}pyodide.mjs`;
const WORKSPACE = '/workspace';
const encoder = new TextEncoder();

type PyodideFs = {
  mkdirTree(path: string): void;
  writeFile(path: string, data: Uint8Array): void;
  readFile(path: string, options?: { encoding?: 'binary' }): Uint8Array;
  readdir(path: string): string[];
  stat(path: string): { mode: number; size: number };
  isDir(mode: number): boolean;
  unlink(path: string): void;
  rmdir(path: string): void;
  chdir(path: string): void;
  cwd(): string;
};

type Pyodide = {
  FS: PyodideFs;
  globals: { get(name: string): unknown };
  loadPackagesFromImports(code: string, options?: { messageCallback?: (message: string) => void; errorCallback?: (message: string) => void }): Promise<void>;
  runPythonAsync(code: string): Promise<unknown>;
  setStdout(options: { batched: (text: string) => void }): void;
  setStderr(options: { batched: (text: string) => void }): void;
  setInterruptBuffer(buffer: Int32Array): void;
  version: string;
  loadedPackages?: Record<string, string>;
};

type PyodideModule = {
  loadPyodide(options?: { indexURL?: string }): Promise<Pyodide>;
};

type WorkerRequest =
  | { id: string; type: 'initialize'; interruptBuffer?: SharedArrayBuffer }
  | { id: string; type: 'workspace-reset' }
  | { id: string; type: 'workspace-mkdir'; path: string }
  | { id: string; type: 'workspace-write'; path: string; data: Uint8Array }
  | { id: string; type: 'workspace-read'; path: string }
  | { id: string; type: 'workspace-list' }
  | { id: string; type: 'execute'; request: RuntimeExecutionRequest }
  | { id: string; type: 'environment' };

type WorkerResponse =
  | { id: string; ok: true; value?: unknown }
  | { id: string; ok: false; error: string };

const scope = globalThis as unknown as {
  postMessage(message: WorkerResponse): void;
  onmessage: ((event: MessageEvent<WorkerRequest>) => void) | null;
};

let pyodidePromise: Promise<Pyodide> | undefined;
let replInitialized = false;

function safeRelativePath(input: string): string {
  const normalized = input.replace(/^\/workspace\/?/u, '').replace(/^\/+/, '');
  const parts: string[] = [];
  for (const segment of normalized.split('/')) {
    if (!segment || segment === '.') continue;
    if (segment === '..') throw new Error('Scientific Python paths are restricted to /workspace');
    parts.push(segment);
  }
  return parts.join('/');
}

function absolutePath(input: string): string {
  const relative = safeRelativePath(input);
  return relative ? `${WORKSPACE}/${relative}` : WORKSPACE;
}

async function pyodide(): Promise<Pyodide> {
  pyodidePromise ??= (async () => {
    const module = await import(/* @vite-ignore */ PYODIDE_MODULE_URL) as PyodideModule;
    const runtime = await module.loadPyodide({ indexURL: PYODIDE_INDEX_URL });
    try { runtime.FS.mkdirTree(WORKSPACE); } catch { /* already exists */ }
    runtime.FS.chdir(WORKSPACE);
    // The worker has no DOM. Use Matplotlib's file-backed backend so research
    // scripts can reproducibly save PDF/SVG/PNG into /workspace/figures.
    await runtime.runPythonAsync("import os; os.environ.setdefault('MPLBACKEND', 'Agg')");
    return runtime;
  })().catch((error) => {
    pyodidePromise = undefined;
    throw error;
  });
  return pyodidePromise;
}

function removeTree(fs: PyodideFs, path: string): void {
  let names: string[];
  try { names = fs.readdir(path); } catch { return; }
  for (const name of names) {
    if (name === '.' || name === '..') continue;
    const child = `${path}/${name}`;
    const stat = fs.stat(child);
    if (fs.isDir(stat.mode)) {
      removeTree(fs, child);
      fs.rmdir(child);
    } else {
      fs.unlink(child);
    }
  }
}

function listFilesRecursive(fs: PyodideFs, root = WORKSPACE): RuntimeWorkspaceFile[] {
  const result: RuntimeWorkspaceFile[] = [];
  const visit = (directory: string) => {
    for (const name of fs.readdir(directory)) {
      if (name === '.' || name === '..') continue;
      const full = `${directory}/${name}`;
      const stat = fs.stat(full);
      if (fs.isDir(stat.mode)) visit(full);
      else result.push({ path: full.slice(`${WORKSPACE}/`.length), size: stat.size });
    }
  };
  visit(root);
  return result.sort((a, b) => a.path.localeCompare(b.path));
}

function pythonString(value: string): string {
  return JSON.stringify(value);
}

async function execute(runtime: Pyodide, request: RuntimeExecutionRequest): Promise<RuntimeExecutionResult> {
  if (request.mode === 'tool') throw new Error(`Python runtime does not expose structured tool: ${request.tool}`);
  const startedAt = Date.now();
  const stdout: string[] = [];
  const stderr: string[] = [];
  runtime.setStdout({ batched: (text) => stdout.push(`${text}\n`) });
  runtime.setStderr({ batched: (text) => stderr.push(`${text}\n`) });
  const cwd = absolutePath(request.cwd ?? '');
  runtime.FS.chdir(cwd);

  let userCode = '';
  if (request.mode === 'file') {
    const sourcePath = absolutePath(request.path);
    userCode = new TextDecoder().decode(runtime.FS.readFile(sourcePath, { encoding: 'binary' }));
  } else {
    userCode = request.code;
  }

  try {
    await runtime.loadPackagesFromImports(userCode, {
      messageCallback: (message) => stdout.push(`[Pyodide] ${message}\n`),
      errorCallback: (message) => stderr.push(`[Pyodide] ${message}\n`),
    });

    const args = request.args ?? [];
    const sourceName = request.mode === 'file'
      ? absolutePath(request.path)
      : request.sourceName ?? '<interactive-python>';
    const isRepl = request.mode === 'code' && request.sourceName === '<interactive-repl>';
    let wrapper: string;
    if (isRepl) {
      if (!replInitialized) {
        await runtime.runPythonAsync('import code as _learnitive_code\n_learnitive_console = _learnitive_code.InteractiveConsole()');
        replInitialized = true;
      }
      wrapper = `_learnitive_console.push(${pythonString(userCode)})`;
    } else if (request.mode === 'file') {
      wrapper = [
        'import os, runpy, sys',
        `os.chdir(${pythonString(cwd)})`,
        `sys.argv = [${pythonString(sourceName)}] + ${JSON.stringify(args)}`,
        `runpy.run_path(${pythonString(sourceName)}, run_name='__main__')`,
      ].join('\n');
    } else {
      wrapper = [
        'import os, sys',
        `os.chdir(${pythonString(cwd)})`,
        `sys.argv = [${pythonString(sourceName)}] + ${JSON.stringify(args)}`,
        `_learnitive_globals = {'__name__': '__main__', '__file__': ${pythonString(sourceName)}}`,
        `exec(compile(${pythonString(userCode)}, ${pythonString(sourceName)}, 'exec'), _learnitive_globals, _learnitive_globals)`,
      ].join('\n');
    }
    await runtime.runPythonAsync(wrapper);
    return {
      runtime: 'python',
      exitCode: 0,
      stdout: stdout.join(''),
      stderr: stderr.join(''),
      startedAt,
      durationMs: Date.now() - startedAt,
      metadata: {
        pyodideVersion: PYODIDE_VERSION,
        execution: isRepl ? 'repl' : request.mode,
        loadedPackages: Object.keys(runtime.loadedPackages ?? {}).join(','),
      },
    };
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    if (!stderr.join('').includes(message)) stderr.push(`${message}\n`);
    const interrupted = /KeyboardInterrupt|Interrupted|AbortError/iu.test(message);
    return {
      runtime: 'python',
      exitCode: interrupted ? 130 : 1,
      stdout: stdout.join(''),
      stderr: stderr.join(''),
      startedAt,
      durationMs: Date.now() - startedAt,
      metadata: {
        pyodideVersion: PYODIDE_VERSION,
        loadedPackages: Object.keys(runtime.loadedPackages ?? {}).join(','),
      },
    };
  } finally {
    try { runtime.FS.chdir(WORKSPACE); } catch { /* runtime may be terminating */ }
  }
}

async function handle(message: WorkerRequest): Promise<unknown> {
  const runtime = await pyodide();
  if (message.type === 'initialize') {
    if (message.interruptBuffer) runtime.setInterruptBuffer(new Int32Array(message.interruptBuffer));
    return { pyodideVersion: PYODIDE_VERSION };
  }
  if (message.type === 'workspace-reset') {
    removeTree(runtime.FS, WORKSPACE);
    runtime.FS.chdir(WORKSPACE);
    replInitialized = false;
    return undefined;
  }
  if (message.type === 'workspace-mkdir') {
    runtime.FS.mkdirTree(absolutePath(message.path));
    return undefined;
  }
  if (message.type === 'workspace-write') {
    const path = absolutePath(message.path);
    const parent = path.slice(0, path.lastIndexOf('/')) || WORKSPACE;
    runtime.FS.mkdirTree(parent);
    runtime.FS.writeFile(path, message.data);
    return undefined;
  }
  if (message.type === 'workspace-read') {
    try {
      const data = runtime.FS.readFile(absolutePath(message.path), { encoding: 'binary' });
      return new Uint8Array(data);
    } catch {
      return undefined;
    }
  }
  if (message.type === 'workspace-list') return listFilesRecursive(runtime.FS);
  if (message.type === 'environment') {
    const version = await runtime.runPythonAsync("import sys; sys.version.split()[0]");
    return { pythonVersion: String(version), pyodideVersion: PYODIDE_VERSION };
  }
  if (message.type === 'execute') return execute(runtime, message.request);
  return undefined;
}

scope.onmessage = (event) => {
  const message = event.data;
  void handle(message).then(
    (value) => scope.postMessage({ id: message.id, ok: true, value }),
    (error) => scope.postMessage({ id: message.id, ok: false, error: error instanceof Error ? error.message : String(error) }),
  );
};