import { serve } from '@hono/node-server'
import { serveStatic } from '@hono/node-server/serve-static'
import { Hono } from 'hono'
import { logger } from 'hono/logger'
import { prettyJSON } from 'hono/pretty-json'
import os from 'node:os'
import { format } from 'node:util'
import { isMainThread, parentPort, Worker } from 'node:worker_threads'
import playwright from 'playwright-extra'
import prettyBytes from 'pretty-bytes'
import prettyMs from 'pretty-ms'
import pluginStealth from 'puppeteer-extra-plugin-stealth'

const TIMEOUT_MS = 6e4
const MAX_CODE_LENGTH = 6e4

if (!isMainThread) {
	const AsyncFunction = Object.getPrototypeOf(async () => {}).constructor
	parentPort.on('message', async (code) => {
		try {
			const fn = new AsyncFunction(
				'playwright',
				'pluginStealth',
				'console',
				code
			)
			parentPort.postMessage({
				result: await fn(playwright, pluginStealth, console)
			})
		} catch (e) {
			parentPort.postMessage({ error: format(e) })
		}
	})
} else {
	const isNumber = (v) => typeof v === 'number' && !isNaN(v)

	const transformObj = (obj, cb) =>
		JSON.parse(
			JSON.stringify(obj, (key, val) =>
				cb(key, val) ? prettyBytes(val) : val
			).replace(/_(\w)/g, (_, g) => g.toUpperCase())
		)

	const getServerStats = () => {
		const stats = {}
		stats.uptime = prettyMs(process.uptime() * 1e3)
		stats.osUptime = prettyMs(os.uptime() * 1e3)

		const report = process.report?.getReport?.()
		if (report)
			Object.assign(stats, {
				header: report.header,
				javascriptHeap: transformObj(
					report.javascriptHeap,
					(k, v) => !/(ContextCount|Garbage)/i.test(k) && isNumber(v)
				),
				resourceUsage: transformObj(
					report.resourceUsage,
					(k, v) =>
						!/(Percent|IO|reads|write)/i.test(k) && isNumber(v)
				),
				uvthreadResourceUsage: transformObj(
					report.uvthreadResourceUsage,
					(k, v) => isNumber(v)
				)
			})

		stats.memoryUsage = transformObj(process.memoryUsage(), (k, v) =>
			isNumber(v)
		)

		return stats
	}

	const runBrowserScript = (code) =>
		new Promise((resolve, reject) => {
			const worker = new Worker(new URL(import.meta.url), {
				resourceLimits: { maxOldGenerationSizeMb: 512 }
			})

			const timer = setTimeout(() => {
				worker.terminate()
				reject(new Error('Execution timeout'))
			}, TIMEOUT_MS)

			const cleanup = () => {
				clearTimeout(timer)
				worker.terminate()
			}

			worker.postMessage(code)
			worker.on('message', ({ result, error }) => {
				cleanup()
				error ? reject(new Error(error)) : resolve(format(result))
			})
			worker.on('error', (e) => {
				cleanup()
				reject(e)
			})
			worker.on('exit', (exitCode) => {
				clearTimeout(timer)
				exitCode !== 0 &&
					reject(new Error(`Worker exited with code ${exitCode}`))
			})
		})

	const app = new Hono()

	app.use(logger())
	app.use(prettyJSON({ force: true }))
	app.use(
		'/file/*',
		serveStatic({
			root: os.tmpdir(),
			rewriteRequestPath: (path) => path.replace(/^\/file/, '')
		})
	)

	app.get('/', (c) => c.json(getServerStats()))

	app.post('/run', async (c) => {
		const body = await c.req.json().catch(() => ({}))
		const { code } = body

		if (!code) return c.json({ error: 'Code is required' }, 400)
		if (typeof code !== 'string')
			return c.json({ error: 'Code must be a string' }, 400)
		if (code.length > MAX_CODE_LENGTH)
			return c.json({ error: 'Code too long' }, 400)

		try {
			return c.json({ result: await runBrowserScript(code) })
		} catch (e) {
			return c.json({ error: format(e) }, 500)
		}
	})

	const port = process.env.SPACE_ID ? 7860 : +(process.env.PORT || 3000)
	serve({ fetch: app.fetch, port }, console.log)
}