import http from 'node:http'
import net from 'node:net'
import { GPU_WAIT_MESSAGE } from './gpu-reservation.mjs'

function reply(res, status, body) {
  if (res.headersSent) return res.destroy()
  res.writeHead(status, { 'Content-Type': 'application/json', 'Cache-Control': 'no-store' })
  res.end(JSON.stringify(body))
}

/** Streams Comfy HTTP/WebSocket traffic; all mutations require the current reservation. */
export function createGpuProxy({ target, reservation, authorized = () => true, markWork = () => {}, externalBusy = () => false }) {
  const server = http.createServer(async (req, res) => {
    const path = new URL(req.url || '/', 'http://localhost').pathname.replace(/\/+$/, '') || '/'
    // A global history list can expose another studio's prompts and outputs.
    if (req.method === 'GET' && (path === '/history' || path === '/history/')) return reply(res, 200, {})
    const forward = () => new Promise(resolve => {
      const port = target()
      if (!port) { reply(res, 503, { error: 'GPU service is unavailable.' }); return resolve() }
      const headers = { ...req.headers }
      delete headers['x-aigen-gpu-lease']
      delete headers.authorization
      const upstream = http.request({ hostname: '127.0.0.1', port, path: req.url, method: req.method, headers }, response => {
        if (req.method === 'GET' && path === '/queue') {
          let size = 0
          const chunks = []
          response.on('data', chunk => {
            size += chunk.length
            if (size > 16 * 1024 * 1024) { response.destroy(); reply(res, 502, { error: 'GPU status unavailable.' }); resolve(); return }
            chunks.push(chunk)
          })
          response.on('end', () => {
            try {
              if (response.statusCode !== 200) throw new Error('status unavailable')
              const queue = JSON.parse(Buffer.concat(chunks).toString('utf8'))
              if (!Array.isArray(queue.queue_running) || !Array.isArray(queue.queue_pending)) throw new Error('invalid queue')
              // Preserve count compatibility without revealing IDs, graphs, or prompts.
              reply(res, 200, { queue_running: queue.queue_running.map(() => null), queue_pending: queue.queue_pending.map(() => null) })
            } catch { reply(res, 502, { error: 'GPU status unavailable.' }) }
            resolve()
          })
        } else {
          res.writeHead(response.statusCode || 502, response.headers)
          response.pipe(res)
          response.on('end', resolve)
        }
        response.on('error', () => { res.destroy(); resolve() })
      })
      upstream.on('error', () => { reply(res, 502, { error: 'GPU service is unavailable.' }); resolve() })
      upstream.setTimeout(5 * 60 * 1000, () => upstream.destroy())
      req.on('aborted', () => upstream.destroy())
      res.on('close', () => { upstream.destroy(); resolve() })
      req.pipe(upstream)
    })
    try {
      if (req.method === 'GET' || req.method === 'HEAD') await forward()
      else {
        if (externalBusy()) { req.resume(); return reply(res, 409, { error: { message: 'A music or upscale job is using the GPU.' }, code: 'GPU_BUSY' }) }
        if (!authorized(req)) { req.resume(); return reply(res, 401, { error: 'Unauthorized' }) }
        await reservation.permit(String(req.headers['x-aigen-gpu-lease'] || ''), async () => { markWork(); await forward() })
      }
    } catch {
      req.resume()
      reply(res, 409, { error: { message: GPU_WAIT_MESSAGE }, code: 'GPU_BUSY' })
    }
  })
  server.on('upgrade', (req, socket, head) => {
    if (new URL(req.url || '/', 'http://localhost').pathname !== '/ws' || !target()) return socket.destroy()
    const upstream = net.connect(target(), '127.0.0.1', () => {
      const headers = Object.entries(req.headers).filter(([key]) => key !== 'authorization' && key !== 'x-aigen-gpu-lease')
        .map(([key, value]) => `${key}: ${value}`).join('\r\n')
      upstream.write(`${req.method} ${req.url} HTTP/${req.httpVersion}\r\n${headers}\r\n\r\n`)
      if (head.length) upstream.write(head)
      socket.pipe(upstream); upstream.pipe(socket)
    })
    socket.on('error', () => upstream.destroy())
    upstream.on('error', () => socket.destroy())
    socket.on('close', () => upstream.destroy())
    upstream.on('close', () => socket.destroy())
  })
  return server
}