import http, { type IncomingHttpHeaders, type IncomingMessage, type ServerResponse } from 'node:http'
import https from 'node:https'
import { once } from 'node:events'
import type { Config, Rule, Upstream } from './config.js'

const STRIP_REQUEST_HEADERS = new Set([
  'authorization', 'cookie', 'host', 'proxy-authorization', 'proxy-authenticate', 'forwarded', 'x-real-ip',
  'connection', 'keep-alive', 'te', 'trailer', 'transfer-encoding', 'upgrade',
])
const STRIP_RESPONSE_HEADERS = new Set([
  'set-cookie', 'set-cookie2', 'proxy-authenticate', 'connection', 'keep-alive', 'te', 'trailer',
  'transfer-encoding', 'upgrade',
])

function canonicalPath(pathname: string): string {
  let decoded = pathname
  for (let i = 0; i < 4; i += 1) {
    const next = decodeURIComponent(decoded)
    if (next === decoded) break
    decoded = next
  }
  if (/%[0-9A-Fa-f]{2}/.test(decoded) || decoded.includes('\\') || decoded.includes('\0')) throw new Error('invalid encoded path')
  return new URL(decoded, 'http://keyhole.invalid').pathname
}

function pathMatches(rules: Rule[], path: string): boolean {
  const candidate = path.toLowerCase()
  return rules.some((rule) => {
    const selected = canonicalPath(rule.path).toLowerCase()
    return rule.subtree
      ? selected === '/' || candidate === selected || candidate.startsWith(selected.endsWith('/') ? selected : `${selected}/`)
      : candidate === selected
  })
}

async function readBody(request: IncomingMessage, limit: number): Promise<Buffer> {
  const chunks: Buffer[] = []
  let length = 0
  for await (const chunk of request) {
    const bytes = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)
    length += bytes.length
    if (length > limit) throw Object.assign(new Error('request body too large'), { statusCode: 413 })
    chunks.push(bytes)
  }
  return Buffer.concat(chunks)
}

function copyHeaders(headers: IncomingHttpHeaders): http.OutgoingHttpHeaders {
  const result: http.OutgoingHttpHeaders = {}
  for (const [name, value] of Object.entries(headers)) {
    const lower = name.toLowerCase()
    if (!STRIP_REQUEST_HEADERS.has(lower) && !lower.startsWith('x-forwarded-') && value !== undefined) result[lower] = value
  }
  return result
}

function requestUpstream(method: string, target: URL, headers: http.OutgoingHttpHeaders, body: Buffer): Promise<IncomingMessage> {
  const transport = target.protocol === 'https:' ? https : http
  return new Promise((resolve, reject) => {
    const request = transport.request(target, { method, headers, agent: false, timeout: 30_000 }, resolve)
    request.on('timeout', () => request.destroy(new Error('upstream timeout')))
    request.on('error', reject)
    if (body.length) request.write(body)
    request.end()
  })
}

async function drain(response: IncomingMessage): Promise<void> {
  response.resume()
  await once(response, 'end')
}

function parseAllowedCookies(setCookieHeaders: string[] | string | undefined, names: string[]): string {
  const allowed = new Set(names)
  const cookies: string[] = []
  for (const header of setCookieHeaders === undefined ? [] : Array.isArray(setCookieHeaders) ? setCookieHeaders : [setCookieHeaders]) {
    const pair = header.split(';', 1)[0]
    if (pair === undefined) continue
    const index = pair.indexOf('=')
    if (index > 0 && allowed.has(pair.slice(0, index))) cookies.push(pair)
  }
  if (new Set(cookies.map((pair) => pair.slice(0, pair.indexOf('=')))).size !== allowed.size) {
    throw new Error('login response did not set every required cookie')
  }
  return cookies.join('; ')
}

async function login(upstream: Upstream): Promise<string> {
  const auth = upstream.auth
  if (auth.type !== 'cookie-login') throw new Error('cookie login is not configured')
  const target = new URL(auth.loginPath, upstream.baseUrl)
  let body: Buffer
  let contentType: string
  if (auth.bodyType === 'json') {
    body = Buffer.from(JSON.stringify(auth.body))
    contentType = 'application/json'
  } else {
    body = Buffer.from(new URLSearchParams(auth.body).toString())
    contentType = 'application/x-www-form-urlencoded'
  }
  const response = await requestUpstream('POST', target, { 'content-type': contentType, 'content-length': String(body.length) }, body)
  const setCookie = response.headersDistinct?.['set-cookie'] ?? response.headers['set-cookie']
  if ((response.statusCode ?? 500) < 200 || (response.statusCode ?? 500) >= 300) {
    await drain(response)
    throw new Error('upstream login failed')
  }
  const cookie = parseAllowedCookies(setCookie, auth.cookieNames)
  await drain(response)
  return cookie
}

async function authHeaders(upstream: Upstream): Promise<http.OutgoingHttpHeaders> {
  const auth = upstream.auth
  if (auth.type === 'none') return {}
  if (auth.type === 'header') return { [auth.header.toLowerCase()]: auth.value }
  if (!auth.cookiePromise) {
    auth.cookiePromise = login(upstream).catch((error: unknown) => {
      auth.cookiePromise = null
      throw error
    })
  }
  return { cookie: await auth.cookiePromise }
}

async function forward(upstream: Upstream, method: string, target: URL, incomingHeaders: IncomingHttpHeaders, body: Buffer, retry = true): Promise<IncomingMessage> {
  const headers = { ...copyHeaders(incomingHeaders), ...(await authHeaders(upstream)) }
  if (body.length) headers['content-length'] = String(body.length)
  else delete headers['content-length']
  const response = await requestUpstream(method, target, headers, body)
  if (upstream.auth.type === 'cookie-login' && retry && upstream.auth.refreshStatuses.includes(response.statusCode ?? 0)) {
    await drain(response)
    upstream.auth.cookiePromise = null
    return forward(upstream, method, target, incomingHeaders, body, false)
  }
  return response
}

function writeError(response: ServerResponse, statusCode: number, code: string): void {
  const body = Buffer.from(JSON.stringify({ error: code }))
  response.writeHead(statusCode, { 'content-type': 'application/json', 'content-length': body.length })
  response.end(body)
}

function parseRoute(requestUrl: string | undefined): { path: string; search: string } | null {
  if (!requestUrl || !requestUrl.startsWith('/') || requestUrl.startsWith('//')) return null
  const parsed = new URL(requestUrl, 'http://keyhole.invalid')
  return { path: canonicalPath(parsed.pathname), search: parsed.search }
}

export function createHandler(config: Config): http.RequestListener {
  return async function handler(request, response) {
    try {
      const host = request.headers.host?.toLowerCase()
      const match = host?.match(/^([a-z0-9](?:[a-z0-9.-]*[a-z0-9])?)(?::([0-9]{1,5}))?$/)
      const domain = match?.[1]
      const upstream = domain && (!match?.[2] || Number(match[2]) === request.socket.localPort) ? config.domains.get(domain) : null
      if (!upstream) return writeError(response, 404, 'not_found')
      const route = parseRoute(request.url)
      if (!route || !pathMatches(upstream.allow, route.path) || pathMatches(upstream.deny, route.path)) return writeError(response, 404, 'not_found')
      const body = await readBody(request, upstream.maxBodyBytes)
      const target = new URL(route.path + route.search, upstream.baseUrl)
      if (target.origin !== upstream.baseUrl.origin) return writeError(response, 400, 'invalid_target')
      const upstreamResponse = await forward(upstream, request.method ?? 'GET', target, request.headers, body)
      const headers: http.OutgoingHttpHeaders = {}
      for (const [name, value] of Object.entries(upstreamResponse.headers)) {
        if (!STRIP_RESPONSE_HEADERS.has(name.toLowerCase()) && value !== undefined) headers[name] = value
      }
      response.writeHead(upstreamResponse.statusCode ?? 502, headers)
      upstreamResponse.pipe(response)
    } catch (error: unknown) {
      const statusCode = error !== null && typeof error === 'object' && 'statusCode' in error && typeof error.statusCode === 'number'
        ? error.statusCode
        : 502
      if (!response.headersSent) writeError(response, statusCode, 'upstream_failure')
      else response.destroy()
    }
  }
}