import { readPrefixedEnvFlag } from '#cli/core/env'

type WorkerThreadsModule = typeof import('node:worker_threads')
type WorkerThread = InstanceType<WorkerThreadsModule['Worker']>
type ShikiTokenLines = unknown[]
type RuntimeGlobals = typeof globalThis & {
  process?: NodeJS.Process
  Bun?: unknown
}

type WorkerCall = Readonly<{
  id: number
  type: 'codeToTokensBase'
  code: string
  lang: string
  theme: string
}>

type WorkerReply = Readonly<
  | { id: number, ok: true, tokenLines: ShikiTokenLines }
  | { id: number, ok: false, error: { message: string, stack?: string } }
>

type Pending = Readonly<{
  resolve: (v: ShikiTokenLines) => void
  reject: (e: Error) => void
}>

function runtimeGlobals(): RuntimeGlobals {
  return globalThis as RuntimeGlobals
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return value != null && typeof value === 'object' && !Array.isArray(value)
}

function parseWorkerReply(raw: unknown): WorkerReply | null {
  if (!isRecord(raw))
    return null
  const id = typeof raw.id === 'number' ? raw.id : Number(raw.id)
  if (!Number.isFinite(id))
    return null
  if (raw.ok === true) {
    return {
      id,
      ok: true,
      tokenLines: Array.isArray(raw.tokenLines) ? raw.tokenLines : [],
    }
  }
  if (raw.ok !== false)
    return null
  const error = isRecord(raw.error) ? raw.error : null
  const message =
    typeof error?.message === 'string' && error.message
      ? error.message
      : 'Shiki worker error'
  const stack = typeof error?.stack === 'string' ? error.stack : undefined
  return {
    id,
    ok: false,
    error: {
      message,
      ...(stack ? { stack } : {}),
    },
  }
}

let workerThreads: WorkerThreadsModule | null = null
let worker: WorkerThread | null = null
let workerInitFailed = false
let seq = 0
const pendingById = new Map<number, Pending>()

function wantsWorker(): boolean {
  const v = readPrefixedEnvFlag(undefined, 'SHIKI_WORKER')
  if (v === '0' || v === 'false' || v === 'off' || v === 'no')
    return false
  // Default ON in node-like environments, but always allow opt-out.
  return true
}

export function resetShikiWorkerForTests(): void {
  try {
    worker?.terminate?.()
  }
  catch {}
  worker = null
  workerThreads = null
  workerInitFailed = false
  seq = 0
  for (const p of pendingById.values()) {
    try {
      p.reject(new Error('Shiki worker reset'))
    }
    catch {}
  }
  pendingById.clear()
}

export function shikiWorkerSupported(): boolean {
  // Keep it cheap/sync; actual import happens on demand.
  const globals = runtimeGlobals()
  return Boolean(globals.process?.versions?.node) || globals.Bun != null
}

export function getShikiTokenizerKind():
  | 'worker'
  | 'main'
  | 'disabled'
  | 'failed' {
  if (!wantsWorker())
    return 'disabled'
  if (workerInitFailed)
    return 'failed'
  if (worker)
    return 'worker'
  return 'main'
}

async function ensureWorker(): Promise<WorkerThread | null> {
  if (!wantsWorker())
    return null
  if (worker)
    return worker
  if (workerInitFailed)
    return null

  let wt: WorkerThreadsModule
  try {
    wt = workerThreads ?? (await import('node:worker_threads'))
    workerThreads = wt
  }
  catch {
    workerInitFailed = true
    return null
  }

  const src = `
    const { parentPort } = require('node:worker_threads');
    let shikiModP = null;
    async function getCodeToTokensBase() {
      if (!shikiModP) shikiModP = import('shiki');
      const mod = await shikiModP;
      const fn = mod && mod.codeToTokensBase;
      if (typeof fn !== 'function') throw new Error('shiki_unavailable');
      return fn;
    }
    parentPort.on('message', async (msg) => {
      const id = Number(msg && msg.id);
      try {
        const fn = await getCodeToTokensBase();
        const code = String(msg && msg.code || '');
        const lang = String(msg && msg.lang || 'text');
        const theme = String(msg && msg.theme || 'monokai');
        const tokenLines = await fn(code, { lang, theme });
        parentPort.postMessage({ id, ok: true, tokenLines });
      } catch (e) {
        const err = e instanceof Error ? e : new Error(String(e));
        parentPort.postMessage({ id, ok: false, error: { message: err.message, stack: err.stack } });
      }
    });
  `.trim()

  try {
    const w = new wt.Worker(src, { eval: true })
    w.on('message', (raw: unknown) => {
      const reply = parseWorkerReply(raw)
      if (!reply)
        return
      const pending = pendingById.get(reply.id)
      if (!pending)
        return
      pendingById.delete(reply.id)
      if (reply.ok) {
        pending.resolve(reply.tokenLines)
      }
      else {
        const e = new Error(reply.error.message)
        if (reply.error.stack)
          e.stack = reply.error.stack
        pending.reject(e)
      }
    })
    w.on('error', (error) => {
      workerInitFailed = true
      worker = null
      for (const p of pendingById.values()) {
        try {
          p.reject(error)
        }
        catch {}
      }
      pendingById.clear()
    })
    w.on('exit', () => {
      worker = null
    })
    worker = w
    return worker
  }
  catch {
    workerInitFailed = true
    worker = null
    return null
  }
}

export async function codeToTokensBaseOffThread(
  code: string,
  opts: Readonly<{ lang: string, theme: string }>,
): Promise<ShikiTokenLines | null> {
  const w = await ensureWorker()
  if (!w)
    return null
  const id = ++seq
  const payload: WorkerCall = {
    id,
    type: 'codeToTokensBase',
    code: String(code ?? ''),
    lang: String(opts.lang ?? 'text'),
    theme: String(opts.theme ?? 'monokai'),
  }
  return await new Promise<ShikiTokenLines>((resolve, reject) => {
    pendingById.set(id, { resolve, reject })
    try {
      w.postMessage(payload)
    }
    catch (e) {
      pendingById.delete(id)
      reject(e instanceof Error ? e : new Error(String(e)))
    }
  })
}
