// ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

import { loadPyodideRuntime, type PyodideInterface } from './loader.ts'
import {
  createMirageBridge,
  preloadInto,
  type BridgeDispatchFn,
  type MirageBridge,
} from './mirage_bridge.ts'
import { MIRAGE_FS_SHIM_PY } from './mirage_fs_shim.ts'
import { PYTHON_REPL_WRAPPER, PYTHON_WRAPPER } from './wrapper.ts'
import type {
  PythonReplRunArgs,
  PythonReplRunResult,
  PythonRunArgs,
  PythonRunResult,
  ReplStatus,
} from './types.ts'

function runtimeEnv(): Record<string, string> {
  const env: Record<string, string> = {}
  const proc = (globalThis as { process?: { env?: Record<string, string | undefined> } }).process
  if (proc?.env === undefined) return env
  for (const [k, v] of Object.entries(proc.env)) {
    if (typeof v === 'string') env[k] = v
  }
  return env
}

function normalizePrefix(prefix: string): string {
  return prefix.endsWith('/') ? prefix : prefix + '/'
}

/**
 * Rewrite top-level imports of denied packages so Pyodide's
 * `loadPackagesFromImports` skips fetching them. The rewritten code is only
 * fed to the auto-loader's import scanner — user code still runs unchanged,
 * so the actual `import X` will hit any meta_path blocker installed in the
 * Python bootstrap.
 *
 * Recognises:
 *   - `import X`, `import X.Y`, `import X as alias`
 *   - `from X import …`, `from X.Y import …`
 * The match is line-scoped (`/m`) so multi-import lines like
 * `import X, Y` are blanked out as a single statement.
 */
export function stripDeniedImports(code: string, denyPackages: ReadonlySet<string>): string {
  if (denyPackages.size === 0) return code
  return code.replace(
    /^[ \t]*(?:from|import)\s+([\w][\w.]*)[^\n]*/gm,
    (match, mod: string): string => {
      const top = mod.split('.')[0] ?? ''
      if (!denyPackages.has(top)) return match
      return match.replace(mod, 'os')
    },
  )
}

export interface PyodideRuntimeOptions {
  autoLoadFromImports?: boolean
  bootstrapCode?: string
  workspaceBridge?: BridgeDispatchFn
  denyPackages?: readonly string[]
}

export class PyodideRuntime {
  private pyodide: PyodideInterface | null = null
  private initPromise: Promise<PyodideInterface> | null = null
  private bootstrapPromise: Promise<void> | null = null
  private queue: Promise<unknown> = Promise.resolve()
  private readonly autoLoadFromImports: boolean
  private readonly bootstrapCode: string | null
  private readonly workspaceBridge: BridgeDispatchFn | null
  private readonly denyPackages: ReadonlySet<string>
  private bridge: MirageBridge | null = null
  private readonly mountedPrefixes = new Set<string>()

  constructor(options: PyodideRuntimeOptions = {}) {
    this.autoLoadFromImports = options.autoLoadFromImports ?? true
    this.bootstrapCode = options.bootstrapCode ?? null
    this.workspaceBridge = options.workspaceBridge ?? null
    this.denyPackages = new Set(options.denyPackages ?? [])
  }

  async addMount(prefix: string): Promise<void> {
    if (this.workspaceBridge === null) {
      throw new Error('PyodideRuntime: addMount requires workspaceBridge option')
    }
    const norm = normalizePrefix(prefix)
    const task = (): Promise<void> => this.addMountOne(norm)
    const next = this.queue.then(task, task)
    this.queue = next.catch(() => undefined)
    return next
  }

  async removeMount(prefix: string): Promise<void> {
    if (this.workspaceBridge === null) {
      throw new Error('PyodideRuntime: removeMount requires workspaceBridge option')
    }
    const norm = normalizePrefix(prefix)
    const task = (): Promise<void> => this.removeMountOne(norm)
    const next = this.queue.then(task, task)
    this.queue = next.catch(() => undefined)
    return next
  }

  private async addMountOne(prefix: string): Promise<void> {
    if (this.mountedPrefixes.has(prefix)) return
    if (this.pyodide === null && this.initPromise === null) {
      this.mountedPrefixes.add(prefix)
      return
    }
    const pyodide = await this.ensureLoaded()
    if (this.bridge === null) {
      throw new Error('PyodideRuntime: bridge missing after ensureLoaded')
    }
    await preloadInto(pyodide.FS, this.bridge, prefix)
    await pyodide.runPythonAsync(
      `import _mirage_fs_shim; _mirage_fs_shim.register(${JSON.stringify(prefix)})`,
    )
    this.mountedPrefixes.add(prefix)
  }

  private async removeMountOne(prefix: string): Promise<void> {
    if (!this.mountedPrefixes.has(prefix)) return
    if (this.pyodide === null) {
      this.mountedPrefixes.delete(prefix)
      return
    }
    await this.pyodide.runPythonAsync(
      `import _mirage_fs_shim; _mirage_fs_shim.unregister(${JSON.stringify(prefix)})`,
    )
    this.mountedPrefixes.delete(prefix)
  }

  async run(args: PythonRunArgs): Promise<PythonRunResult> {
    const task = (): Promise<PythonRunResult> => this.runOne(args)
    const next = this.queue.then(task, task)
    this.queue = next.catch(() => undefined)
    return next
  }

  async runRepl(args: PythonReplRunArgs): Promise<PythonReplRunResult> {
    const task = (): Promise<PythonReplRunResult> => this.runOneRepl(args)
    const next = this.queue.then(task, task)
    this.queue = next.catch(() => undefined)
    return next
  }

  async close(): Promise<void> {
    try {
      await this.queue
    } catch {
      // queue failures already surfaced to individual callers; safe to swallow here
    }
    this.pyodide = null
    this.initPromise = null
    this.bridge = null
    this.mountedPrefixes.clear()
  }

  private async ensureLoaded(): Promise<PyodideInterface> {
    if (this.pyodide !== null) {
      if (this.bootstrapPromise !== null) await this.bootstrapPromise
      await this.wireBridgeIfNeeded(this.pyodide)
      return this.pyodide
    }
    this.initPromise ??= loadPyodideRuntime()
    this.pyodide = await this.initPromise
    if (this.bootstrapCode !== null) {
      const code = this.bootstrapCode
      const py = this.pyodide
      this.bootstrapPromise = (async () => {
        if (py.loadPackagesFromImports !== undefined) {
          try {
            await py.loadPackagesFromImports(code, { messageCallback: () => undefined })
          } catch {
            // best-effort
          }
        }
        await py.runPythonAsync(code)
      })()
      await this.bootstrapPromise
    }
    await this.wireBridgeIfNeeded(this.pyodide)
    return this.pyodide
  }

  private async wireBridgeIfNeeded(pyodide: PyodideInterface): Promise<void> {
    if (this.workspaceBridge === null || this.bridge !== null) return
    const bridge = createMirageBridge(this.workspaceBridge)
    pyodide.registerJsModule('_mirage_bridge', bridge)
    await pyodide.runPythonAsync(MIRAGE_FS_SHIM_PY)
    this.bridge = bridge
    for (const prefix of this.mountedPrefixes) {
      await preloadInto(pyodide.FS, bridge, prefix)
      await pyodide.runPythonAsync(
        `import _mirage_fs_shim; _mirage_fs_shim.register(${JSON.stringify(prefix)})`,
      )
    }
  }

  private async loadImports(pyodide: PyodideInterface, code: string): Promise<void> {
    if (!this.autoLoadFromImports) return
    if (pyodide.loadPackagesFromImports === undefined) return
    const filtered = stripDeniedImports(code, this.denyPackages)
    try {
      await pyodide.loadPackagesFromImports(filtered, { messageCallback: () => undefined })
    } catch {
      // best-effort: missing/unknown packages will surface as ImportError in user code
    }
  }

  private async runOne(args: PythonRunArgs): Promise<PythonRunResult> {
    const pyodide = await this.ensureLoaded()
    await this.loadImports(pyodide, args.code)
    const mergedEnv = { ...runtimeEnv(), ...args.env }
    const argv = ['-c', ...args.args]
    const stdinBytes = args.stdin ?? new Uint8Array()

    const mergedEnvPy = pyodide.toPy(mergedEnv)
    const argvPy = pyodide.toPy(argv)
    const userGlobalsPy = pyodide.toPy({})

    pyodide.globals.set('_user_code', args.code)
    pyodide.globals.set('_argv', argvPy)
    pyodide.globals.set('_merged_env', mergedEnvPy)
    pyodide.globals.set('_stdin_bytes', stdinBytes)
    pyodide.globals.set('_user_globals', userGlobalsPy)

    try {
      await pyodide.runPythonAsync(PYTHON_WRAPPER)
      const resultProxy = pyodide.globals.get('_result') as
        | {
            toJs?: (opts?: Record<string, unknown>) => unknown
            destroy?: () => void
          }
        | null
        | undefined
      const arr = resultProxy?.toJs?.({ create_proxies: false }) as
        | [Uint8Array, Uint8Array, number]
        | undefined
      resultProxy?.destroy?.()
      if (arr === undefined) {
        return {
          stdout: new Uint8Array(),
          stderr: new TextEncoder().encode('python3: runtime returned no result\n'),
          exitCode: 1,
        }
      }
      return {
        stdout: arr[0] instanceof Uint8Array ? arr[0] : new Uint8Array(arr[0] as ArrayLike<number>),
        stderr: arr[1] instanceof Uint8Array ? arr[1] : new Uint8Array(arr[1] as ArrayLike<number>),
        exitCode: arr[2],
      }
    } finally {
      pyodide.globals.delete?.('_user_code')
      pyodide.globals.delete?.('_argv')
      pyodide.globals.delete?.('_merged_env')
      pyodide.globals.delete?.('_stdin_bytes')
      pyodide.globals.delete?.('_user_globals')
      pyodide.globals.delete?.('_result')
      const maybeDestroy = (obj: unknown): void => {
        if (obj !== null && typeof obj === 'object' && 'destroy' in obj) {
          try {
            ;(obj as { destroy: () => void }).destroy()
          } catch {
            // destroy is best-effort; ignore double-destroy errors
          }
        }
      }
      maybeDestroy(mergedEnvPy)
      maybeDestroy(argvPy)
      maybeDestroy(userGlobalsPy)
    }
  }

  private async runOneRepl(args: PythonReplRunArgs): Promise<PythonReplRunResult> {
    const pyodide = await this.ensureLoaded()
    await this.loadImports(pyodide, args.code)

    pyodide.globals.set('_user_code', args.code)
    pyodide.globals.set('_repl_session_id', args.sessionId)

    try {
      await pyodide.runPythonAsync(PYTHON_REPL_WRAPPER)
      const resultProxy = pyodide.globals.get('_repl_result') as
        | {
            toJs?: (opts?: Record<string, unknown>) => unknown
            destroy?: () => void
          }
        | null
        | undefined
      const arr = resultProxy?.toJs?.({ create_proxies: false }) as
        | [Uint8Array, Uint8Array, number, ReplStatus]
        | undefined
      resultProxy?.destroy?.()
      if (arr === undefined) {
        return {
          stdout: new Uint8Array(),
          stderr: new TextEncoder().encode('python3: repl returned no result\n'),
          exitCode: 1,
          status: 'complete',
        }
      }
      return {
        stdout: arr[0] instanceof Uint8Array ? arr[0] : new Uint8Array(arr[0] as ArrayLike<number>),
        stderr: arr[1] instanceof Uint8Array ? arr[1] : new Uint8Array(arr[1] as ArrayLike<number>),
        exitCode: arr[2],
        status: arr[3],
      }
    } finally {
      pyodide.globals.delete?.('_user_code')
      pyodide.globals.delete?.('_repl_session_id')
      pyodide.globals.delete?.('_repl_result')
    }
  }
}
