package com.jihe.punchnet.sdlan.network

import com.google.protobuf.kotlin.toByteString
import com.jihe.punchnet.protobuf.PunchProto.SDLData
import com.jihe.punchnet.sdlan.config.EtherType
import com.jihe.punchnet.sdlan.logs.TerminalLogger
import com.jihe.punchnet.sdlan.utils.ipToString
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.async
import kotlinx.coroutines.delay
import java.util.concurrent.ConcurrentHashMap

const val BroadcastIP = 0xFFFFFFFF
@OptIn(ExperimentalUnsignedTypes::class)
val BroadcastMac = ubyteArrayOf(0xffu, 0xffu, 0xffu, 0xffu, 0xffu, 0xffu).toByteArray()

class ARPTable(
    val scope: CoroutineScope,
) {
    val routeTable: RouteTable = RouteTable()

    val content = ConcurrentHashMap<Int, ARPInfo>()

    suspend fun agingARP() {
        // CoroutineScope(Dispatchers.Default).async {
        scope.async {
            while(true) {
                delay(20_000)
                val now = System.currentTimeMillis()/1000
                val toRemove = mutableListOf<Int>()
                for ((k, v) in content) {
                    if (now-v.timestamp > EtherType.ARP_EXPIRE_TIME) {
                        toRemove.add(k)
                    }
                }

                for (item in toRemove) {
                    content.remove(item)
                }
            }
        }
    }

    fun getMacFromIP(ip: Int): Pair<ByteArray?, Int> {
        val gw = routeTable.getGateway(ip)
        if (gw == null) {
            // not found in route table, just use the ip
            val value = content.get(ip)
            return Pair(value?.mac, ip)
        } else {
            // gw not null, try find the gw's mac
            val value = content.get(gw)
            return Pair(value?.mac, gw)
        }
    }

    fun addToARPTable(ip: Int, mac: ByteArray) {
        val origin = content.get(ip)
        val now = System.currentTimeMillis()/1000
        if (origin != null) {
            TerminalLogger.debugf { "adding ${ipToString(ip)} to arptable" }
            origin.timestamp = now
            origin.mac = mac
            return
        }
        if (content.size >= EtherType.ARP_TABLE_SIZE) {
            // loop and clear the old aged item, or if no one expires,
            // remove the oldest one
            var maxage = -1L

            val toRemove = mutableListOf<Int>()
            var maxindex = -1
            val expireage = EtherType.ARP_EXPIRE_TIME + EtherType.ARP_DETECT_TIMES*5
            for ((k, v) in content) {
                if (now-v.timestamp > expireage) {
                    toRemove.add(k)
                    continue
                }
                if (now-v.timestamp > maxage) {
                    maxage = now-v.timestamp
                    maxindex = k
                }
            }
            if (toRemove.size == 0 && maxindex != -1) {
                toRemove.add(maxindex)
            }

            for (item in toRemove) {
                content.remove(item)
            }
            println("adding ${ipToString(ip)}")
        }
        content[ip] = ARPInfo(now, mac)
    }
}

class ARPInfo(var timestamp: Long, var mac: ByteArray)


class ARPWaitList {
    val content = ConcurrentHashMap<Int, MutableList<ARPWaitInfo>>()

    fun addToWaitList(ip: Int, originData: ByteArray) {
        val origin = content.getOrPut(ip, {mutableListOf<ARPWaitInfo>()})

        val now = System.currentTimeMillis()/1000
        if (origin.size > 5) {
            origin.removeAt(0)
            // origin.removeFirst()
        }
        origin.add(ARPWaitInfo(now, originData))
    }

    suspend fun arpArrived(node: Node, ip: Int, mac: Mac) {
        val waitlist = content.remove(ip)
        if (waitlist == null) {
            return
        }
        if (!node.aes.isAuthorized()) {
            return
        }

        val now = System.currentTimeMillis()/1000
        val networkid = node.networkID.get()
        for (item in waitlist) {
            if (now - item.timestamp > 5) {
                // just skip the packet
                continue
            }
            val packet = formEthernetPacket(node.mac.toByteArray(), mac.toByteArray(), item.originData)
            val size = packet.remaining()
            val encrypted = node.aes.encrypt(packet)
            if (encrypted != null) {
                val data = SDLData.newBuilder()
                    .setIsP2P(true)
                    .setNetworkId(networkid)
                    .setTtl(2)
                    .setSrcMac(node.mac)
                    .setDstMac(mac)
                    .setData(encrypted.toByteString())
                    .build()
                val msg = encodeToUDPMessage(data, PacketType.Data)
                sendPacketToNet(node, mac, msg, size.toLong())
            }

        }
    }
}
class ARPWaitInfo(val timestamp: Long, val originData: ByteArray)
