/* eslint-disable no-console */
/***
 * This is a little server that emulates the protocol used by cloudflare's browser rendering API. In
 * local development, you can run this server, and connect to it instead of cloudflare's (strictly
 * limited) API. e.g. in your worker you might use a function like this:
 *
 * ```ts
 * import { Browser, launch as launchPuppeteer } from '@cloudflare/puppeteer'
 * function launchBrowser(env: Environment) {
 * 	 if (env.LOCAL_BROWSER_ORIGIN) {
 * 	   return launchPuppeteer({
 *       fetch: (input, init) => {
 *         const request = new Request(input, init)
 *         return fetch(`${LOCAL_BROWSER_ORIGIN}${request.url}`, request)
 *       },
 *     })
 *   } else {
 *     return launchPuppeteer(env.BROWSER)
 *   }
 * }
 * ```
 *
 * Run this file with `tsx` - https://www.npmjs.com/package/tsx: `tsx runDevBrowser.ts`
 * Or use the JavaScript version (runDevBrowser.mjs) below
 */
import { ServerResponse, createServer } from 'http'
import * as puppeteer from 'puppeteer'
import { WebSocket, WebSocketServer } from 'ws'

const browsers = new Map<string, { lastUsed: number; instance: puppeteer.Browser }>()

setInterval(async () => {
	for (const [id, browser] of browsers) {
		if (browser.lastUsed < Date.now() - 10000) {
			console.log(`closing browser session: ${id}`)
			await browser.instance.close()
			browsers.delete(id)
		}
	}
}, 1000)

const httpServer = createServer(async (req, res) => {
	try {
		console.log(req.method, req.url)
		const [path] = req.url!.split('?')
		switch (path) {
			case '/v1/acquire': {
				const id = `local-${Math.random().toString(36).slice(2)}`
				const browser = await puppeteer.launch({ headless: 'new' })
				console.log(`launching browser session: ${id}`)
				browsers.set(id, { lastUsed: Date.now(), instance: browser })
				sendJson(res, { sessionId: id })
				return
			}
			case '/v1/connectDevtools':
				return
			default:
				console.log('not found:', req.url)
				res.statusCode = 404
				res.end('not found')
				return
		}
	} catch (err) {
		console.log(err)
		res.statusCode = 500
		res.end('error')
	}
})

const wsServer = new WebSocketServer({ noServer: true })
httpServer.on('upgrade', (req, socketToWorker, head) => {
	console.log('UPGRADE', req.url)
	const [path, search] = req.url!.split('?')
	if (path !== '/v1/connectDevtools') {
		socketToWorker.destroy()
		return
	}
	const searchParams = new URLSearchParams(search)
	const sessionId = searchParams.get('browser_session')

	const browser = browsers.get(sessionId!)
	if (!browser) {
		console.log('browser not found')
		socketToWorker.destroy()
		return
	}

	const browserWsUrl = browser.instance.wsEndpoint()
	const socketToBrowser = new WebSocket(browserWsUrl)
	let _socketToWorker: WebSocket | null = null
	socketToBrowser.on('error', (err) => console.log('browser socket error', err))
	socketToBrowser.on('close', () => {
		console.log('b: close')
		if (_socketToWorker) _socketToWorker.close()
	})
	socketToBrowser.on('open', () => {
		const chunksFromWorker: Uint8Array[] = []

		wsServer.handleUpgrade(req, socketToWorker, head, (socketToWorker) => {
			_socketToWorker = socketToWorker

			socketToWorker.on('error', (err) => console.log('worker socket error', err))
			socketToWorker.on('message', (data) => {
				browser.lastUsed = Date.now()
				if (data.toString('utf8') === 'ping') return
				chunksFromWorker.push(new Uint8Array(data as ArrayBuffer))
				const message = chunking.chunksToMessage(chunksFromWorker, sessionId!)
				if (message) {
					socketToBrowser.send(message)
				}
			})
			socketToWorker.on('close', () => {
				console.log('w: close')
				socketToBrowser.close()
			})

			socketToBrowser.on('message', (data) => {
				const chunks = chunking.messageToChunks(data.toString('utf8'))
				for (const chunk of chunks) {
					socketToWorker.send(chunk)
				}
			})
		})
	})
})

function sendJson(res: ServerResponse, data: unknown) {
	res.setHeader('content-type', 'application/json')
	res.end(JSON.stringify(data))
}

httpServer.listen(8789, () => {
	console.log('Listening on port 8789')
})

/**
 * Chunking - adapted from https://github.com/cloudflare/puppeteer/blob/main/src/common/chunking.ts#L27
 * @license Apache-2.0
 */
const chunking = (() => {
	const HEADER_SIZE = 4 // Uint32
	const MAX_MESSAGE_SIZE = 1048575 // Workers size is < 1MB
	const FIRST_CHUNK_DATA_SIZE = MAX_MESSAGE_SIZE - HEADER_SIZE

	const messageToChunks = (data: string): Uint8Array[] => {
		const encoder = new TextEncoder()
		const encodedUint8Array = encoder.encode(data)

		// We only include the header into the first chunk
		const firstChunk = new Uint8Array(
			Math.min(MAX_MESSAGE_SIZE, HEADER_SIZE + encodedUint8Array.length)
		)
		const view = new DataView(firstChunk.buffer)
		view.setUint32(0, encodedUint8Array.length, true)
		firstChunk.set(encodedUint8Array.slice(0, FIRST_CHUNK_DATA_SIZE), HEADER_SIZE)

		const chunks: Uint8Array[] = [firstChunk]
		for (let i = FIRST_CHUNK_DATA_SIZE; i < data.length; i += MAX_MESSAGE_SIZE) {
			chunks.push(encodedUint8Array.slice(i, i + MAX_MESSAGE_SIZE))
		}
		return chunks
	}

	const chunksToMessage = (chunks: Uint8Array[], sessionid: string): string | null => {
		if (chunks.length === 0) {
			return null
		}

		const emptyBuffer = new Uint8Array(0)
		const firstChunk = chunks[0] || emptyBuffer
		const view = new DataView(firstChunk.buffer)
		const expectedBytes = view.getUint32(0, true)

		let totalBytes = -HEADER_SIZE
		for (let i = 0; i < chunks.length; ++i) {
			const curChunk = chunks[i] || emptyBuffer
			totalBytes += curChunk.length

			if (totalBytes > expectedBytes) {
				throw new Error(
					`Should have gotten the exact number of bytes but we got more.  SessionID: ${sessionid}`
				)
			}
			if (totalBytes === expectedBytes) {
				const chunksToCombine = chunks.splice(0, i + 1)
				chunksToCombine[0] = firstChunk.subarray(HEADER_SIZE)

				const combined = new Uint8Array(expectedBytes)
				let offset = 0
				for (let j = 0; j <= i; ++j) {
					const chunk = chunksToCombine[j] || emptyBuffer
					combined.set(chunk, offset)
					offset += chunk.length
				}
				const decoder = new TextDecoder()
				// return decoder.decode(combined)
				const message = decoder.decode(combined)
				return message
			}
		}
		return null
	}

	return { chunksToMessage, messageToChunks }
})()