package expo.modules.rebreakprotection.filter

import android.util.Log
import java.io.FileOutputStream
import java.net.DatagramPacket
import java.net.DatagramSocket
import java.net.InetAddress
import java.util.concurrent.Executors
import java.util.concurrent.ThreadPoolExecutor

/**
 * IPv4 + UDP + DNS Packet-Parsing und Response-Builder.
 *
 * Verarbeitungsmodell:
 *  - Block-Queries → sofort NXDOMAIN-Response (synchron, returns ByteArray)
 *  - Allow-Queries → async Thread-Pool: forwarded an Upstream, schreibt Response
 *    direkt ins TUN (synchronized auf outputLock).
 *
 * Damit blockt EINE langsame Forward-Query nicht den ganzen Filter-Loop.
 */
object DnsFilter {

    private const val TAG = "RebreakDnsFilter"
    private const val UPSTREAM_DNS = "1.1.1.1"
    private const val UPSTREAM_PORT = 53
    private const val DNS_PORT = 53

    private val forwardPool: ThreadPoolExecutor =
        Executors.newCachedThreadPool() as ThreadPoolExecutor

    /**
     * Verarbeitet ein Paket vom TUN.
     * - Bei Block: schreibt synchron NXDOMAIN-Response in `output`.
     * - Bei Allow: fired Thread für Forward + Async-Write.
     * - Bei Drop (kein DNS, IPv6, etc.): nichts.
     *
     * Caller (VpnService) muss vor jedem `output.write(...)` den `outputLock`
     * synchronisieren — wir tun das hier intern.
     */
    fun process(
        packet: ByteArray,
        length: Int,
        hashList: HashList,
        output: FileOutputStream,
        outputLock: Any,
        protectSocket: (DatagramSocket) -> Boolean,
    ) {
        if (length < 28) return // min IPv4(20) + UDP(8)

        // IPv4-Check
        val version = (packet[0].toInt() ushr 4) and 0xF
        if (version != 4) return

        val ihl = (packet[0].toInt() and 0xF) * 4
        if (length < ihl + 8) return

        // UDP only
        val protocol = packet[9].toInt() and 0xFF
        if (protocol != 17) return

        // UDP-Header
        val udpStart = ihl
        val srcPort = ((packet[udpStart].toInt() and 0xFF) shl 8) or (packet[udpStart + 1].toInt() and 0xFF)
        val dstPort = ((packet[udpStart + 2].toInt() and 0xFF) shl 8) or (packet[udpStart + 3].toInt() and 0xFF)
        if (dstPort != DNS_PORT) return

        val dnsStart = udpStart + 8
        val dnsLen = length - dnsStart
        if (dnsLen < 12) return

        val domain = parseQname(packet, dnsStart + 12) ?: return

        val srcIp = packet.copyOfRange(12, 16)
        val dstIp = packet.copyOfRange(16, 20)

        if (hashList.matchesAnySuffix(domain)) {
            Log.i(TAG, "BLOCKED: $domain")
            val resp = buildNxDomainResponse(packet, length, ihl, srcIp, dstIp, srcPort, dstPort)
            writeSynchronized(output, outputLock, resp)
            return
        }

        // Async forward — kopiere Buffer-Slice damit nicht überschrieben wird beim
        // nächsten read() im VpnService-Loop
        val packetCopy = packet.copyOf(length)
        forwardPool.execute {
            forwardAndWrite(
                packetCopy, length, ihl, srcIp, dstIp, srcPort, dstPort,
                output, outputLock, protectSocket, domain,
            )
        }
    }

    private fun parseQname(packet: ByteArray, start: Int): String? {
        val sb = StringBuilder()
        var pos = start
        while (pos < packet.size) {
            val len = packet[pos].toInt() and 0xFF
            if (len == 0) break
            if ((len and 0xC0) != 0) return null
            pos++
            if (pos + len > packet.size) return null
            if (sb.isNotEmpty()) sb.append('.')
            sb.append(String(packet, pos, len, Charsets.US_ASCII))
            pos += len
        }
        return if (sb.isEmpty()) null else sb.toString()
    }

    private fun buildNxDomainResponse(
        request: ByteArray,
        length: Int,
        ihl: Int,
        srcIp: ByteArray,
        dstIp: ByteArray,
        srcPort: Int,
        dstPort: Int,
    ): ByteArray {
        val resp = request.copyOf(length)
        System.arraycopy(dstIp, 0, resp, 12, 4)
        System.arraycopy(srcIp, 0, resp, 16, 4)
        resp[10] = 0; resp[11] = 0

        val udpStart = ihl
        resp[udpStart] = (dstPort ushr 8).toByte()
        resp[udpStart + 1] = (dstPort and 0xFF).toByte()
        resp[udpStart + 2] = (srcPort ushr 8).toByte()
        resp[udpStart + 3] = (srcPort and 0xFF).toByte()
        resp[udpStart + 6] = 0; resp[udpStart + 7] = 0

        val dnsStart = udpStart + 8
        resp[dnsStart + 2] = 0x81.toByte()
        resp[dnsStart + 3] = 0x83.toByte() // RCODE=3 (NXDOMAIN)
        for (i in 6..11) resp[dnsStart + i] = 0

        return recomputeIpChecksum(resp, ihl)
    }

    private fun forwardAndWrite(
        request: ByteArray,
        length: Int,
        ihl: Int,
        srcIp: ByteArray,
        dstIp: ByteArray,
        srcPort: Int,
        dstPort: Int,
        output: FileOutputStream,
        outputLock: Any,
        protectSocket: (DatagramSocket) -> Boolean,
        domainForLog: String,
    ) {
        var sock: DatagramSocket? = null
        try {
            sock = DatagramSocket()
            if (!protectSocket(sock)) {
                Log.w(TAG, "protect() failed for $domainForLog — DNS forward would loop, dropping")
                return
            }
            sock.soTimeout = 4000

            val dnsStart = ihl + 8
            val dnsLen = length - dnsStart
            val query = request.copyOfRange(dnsStart, dnsStart + dnsLen)

            val upstreamAddr = InetAddress.getByName(UPSTREAM_DNS)
            sock.send(DatagramPacket(query, query.size, upstreamAddr, UPSTREAM_PORT))

            val buf = ByteArray(4096)
            val pkt = DatagramPacket(buf, buf.size)
            sock.receive(pkt)

            val totalLen = ihl + 8 + pkt.length
            val resp = ByteArray(totalLen)
            System.arraycopy(request, 0, resp, 0, ihl)
            resp[2] = (totalLen ushr 8).toByte()
            resp[3] = (totalLen and 0xFF).toByte()
            System.arraycopy(dstIp, 0, resp, 12, 4)
            System.arraycopy(srcIp, 0, resp, 16, 4)
            resp[10] = 0; resp[11] = 0

            val udpStart = ihl
            resp[udpStart] = (dstPort ushr 8).toByte()
            resp[udpStart + 1] = (dstPort and 0xFF).toByte()
            resp[udpStart + 2] = (srcPort ushr 8).toByte()
            resp[udpStart + 3] = (srcPort and 0xFF).toByte()
            val udpLen = 8 + pkt.length
            resp[udpStart + 4] = (udpLen ushr 8).toByte()
            resp[udpStart + 5] = (udpLen and 0xFF).toByte()
            resp[udpStart + 6] = 0; resp[udpStart + 7] = 0

            System.arraycopy(pkt.data, 0, resp, ihl + 8, pkt.length)
            recomputeIpChecksum(resp, ihl)

            writeSynchronized(output, outputLock, resp)
        } catch (e: Exception) {
            Log.w(TAG, "forward failed for $domainForLog: ${e.javaClass.simpleName}: ${e.message}")
        } finally {
            try { sock?.close() } catch (_: Exception) {}
        }
    }

    private fun writeSynchronized(output: FileOutputStream, outputLock: Any, data: ByteArray) {
        synchronized(outputLock) {
            try {
                output.write(data)
            } catch (e: Exception) {
                Log.w(TAG, "TUN write failed: ${e.message}")
            }
        }
    }

    private fun recomputeIpChecksum(packet: ByteArray, ihl: Int): ByteArray {
        packet[10] = 0; packet[11] = 0
        var sum = 0L
        for (i in 0 until ihl step 2) {
            val word = ((packet[i].toInt() and 0xFF) shl 8) or (packet[i + 1].toInt() and 0xFF)
            sum += word
        }
        while (sum ushr 16 != 0L) sum = (sum and 0xFFFF) + (sum ushr 16)
        val checksum = (sum.inv() and 0xFFFF).toInt()
        packet[10] = (checksum ushr 8).toByte()
        packet[11] = (checksum and 0xFF).toByte()
        return packet
    }
}
