import { randomUUIDv7 } from "bun"

if (process.argv.length < 4) {
    console.log(`more args.\n${process.argv[0]} <ip> <port>`)
}

const ip = process.argv[2] ?? '0.0.0.0'
const port = Number(process.argv[3])

if (!Number.isInteger(port) || port < 1 || port > 65535) {
    console.log("bad port")
    process.exit(1)
}

type tcpsocket = {
    mode: 'tcp'
    socket?: Bun.Socket
}
type udpsocket = {
    mode: 'udp'
    handler: (buf: Buffer) => void
    die: () => void
    heartbeat?: Timer
}
type socket = tcpsocket | udpsocket
type connection = {
    ip: string
    port: number
    queue: Buffer[]
    remoteId?: string
    lastSeen: number
} & socket

const udp = await Bun.udpSocket({
    socket: {
        data(socket, data, port, ip) {
            const id = Object.values(sessions).findIndex(s =>
                s.mode === 'udp' &&
                s.ip === ip &&
                s.port === port
            )

            if (id === -1) return

            // send to handler
            if (sessions[id] && sessions[id].mode === 'udp') {
                sessions[id].handler(data)
            }
        },
        error(socket, error) {
            // dead.
            for (const sess of Object.values(sessions).filter(e => e.mode === 'udp')) {
                sess.die()
            }
        }
    }
})

const connectionsPerIp = new Map<string, number>()
const sessions: Record<string, connection> = {}

async function establishTCP(ip: string, port: number): Promise<string | false> {
    const id = randomUUIDv7()

    const conn: connection = {
        port,
        ip,
        mode: 'tcp',
        socket: undefined,
        queue: [],
        lastSeen: Date.now()
    }

    const sockPromise = Bun.connect({
        port,
        hostname: ip,
        reusePort: true,
        tls: false,
        socket: {
            data(socket, data) { conn.queue.push(data) },
            close(socket, error) { delete sessions[id] },
            end(socket) { delete sessions[id] },
            error(socket, error) { delete sessions[id] }
        }
    })

    const timeoutPromise: Promise<0> = new Promise((r) => setTimeout(() => r(0), 5000))

    const sock: Bun.Socket | 0 = await Promise.race([sockPromise, timeoutPromise]).catch(err => {
        console.log(err)
        return 0
    })

    if (sock === 0) return false

    conn.socket = sock
    sessions[id] = conn
    return id
}

async function establishUDP(ip: string, port: number): Promise<string> {
    const id = randomUUIDv7()

    const conn: connection = {
        port,
        ip,
        mode: 'udp',
        handler: undefined as any,
        die: undefined as any,
        heartbeat: undefined,
        queue: [],
        lastSeen: Date.now()
    }

    conn.handler = (data) => {
        if (!conn.heartbeat) {
            const str = data.toString()
            const bnI = str.indexOf('\n')
            if (str.startsWith('T') && bnI !== -1) {
                const tm = Number(str.substring(1, bnI))
                if (!isNaN(tm)) {
                    conn.heartbeat = setInterval(() => {
                        udp.send('N\n', conn.port, conn.ip)
                    }, tm * 1000)
                }
                const temp_len = bnI + 1
                if (str.length > temp_len) {
                    data = Buffer.from(str.substring(temp_len))
                } else {
                    return
                }
            }
        }
        conn.queue.push(data)
    }
    conn.die = () => {
        clearInterval(conn.heartbeat)
        conn.heartbeat = undefined
    }

    sessions[id] = conn
    udp.send("N\n", port, ip) // handshake

    return id
}

const eviction_timeout = 60_000

setInterval(() => {
    for (const key of Object.keys(sessions)) {
        const sess = sessions[key]!
        if (sess.lastSeen > Date.now() - eviction_timeout) continue
        // evict
        if (sess.mode === 'tcp') sess.socket?.close()
        else sess.die()
        delete sessions[key]
    }
}, eviction_timeout)

Bun.serve({
    port,
    hostname: ip,
    routes: {
        '/': new Response('this is a broadlayer server.'),
        '/tcp/connect/:ip/:port': async req => {
            const port = parseInt(req.params.port)
            if (isNaN(port)) return new Response('err_badport')
            if (port > 65535 || port < 0) return new Response('err_badport')

            const clientIp = req.headers.get('x-forwarded-for')?.split(',')[0]?.trim() || 'unknown'
            let count = connectionsPerIp.get(clientIp) || 0
            if (count >= 4) {
                return new Response('err_toomany', { status: 429 })
            }
            connectionsPerIp.set(clientIp, count + 1)

            const res = await establishTCP(req.params.ip, port)
            if (res === false) return new Response('err_timeout')
            return new Response('ok_' + res)
        },
        '/udp/connect/:ip/:port': async req => {
            const port = parseInt(req.params.port)
            if (isNaN(port)) return new Response('err_badport')
            if (port > 65535 || port < 0) return new Response('err_badport')

            const clientIp = req.headers.get('x-forwarded-for')?.split(',')[0]?.trim() || 'unknown'
            let count = connectionsPerIp.get(clientIp) || 0
            if (count >= 4) {
                return new Response('err_toomany', { status: 429 })
            }
            connectionsPerIp.set(clientIp, count + 1)

            return new Response('ok_' + await establishUDP(req.params.ip, port))
        },
        '/poll/:id': async req => { // read and consume queue
            const sess = sessions[req.params.id]
            if (!sess) return new Response('err_badid')
            const resp = JSON.stringify(sess.queue)
            sess.queue = []
            sess.lastSeen = Date.now()
            return new Response(resp)
        },
        '/peek/:id': async req => { // read queue without consuming
            const sess = sessions[req.params.id]
            if (!sess) return new Response('err_badid')
            const resp = JSON.stringify(sess.queue)
            sess.lastSeen = Date.now()
            return new Response(resp)
        },
        '/close/:id': async req => { // closes a connection
            const sess = sessions[req.params.id]
            if (!sess) return new Response('err_badid')
            if (sess.mode === 'tcp') sess.socket?.close()

            delete sessions[req.params.id]

            const clientIp = req.headers.get('x-forwarded-for')?.split(',')[0]?.trim() || 'unknown'
            let count = connectionsPerIp.get(clientIp) || 0
            if (count > 0) connectionsPerIp.set(clientIp, count - 1)
            if (count <= 1) connectionsPerIp.delete(clientIp)

            return new Response('ok')
        },
        '/send/:id': {
            POST: async req => {
                const sess = sessions[req.params.id]
                if (!sess) return new Response('err_badid')

                if (!req.body) return new Response('err_empty')
                const buf = await (await req.body.blob())!.arrayBuffer()
                if (sess.mode === 'tcp') sess.socket?.write(buf)
                else udp.send(buf, sess.port, sess.ip)

                return new Response('ok')
            }
        }
    }
})