package com.jihe.punchnet.sdlan.network

import com.google.protobuf.kotlin.toByteString
import com.jihe.punchnet.protobuf.PunchProto.SDLData
import com.jihe.punchnet.protobuf.PunchProto.SDLPeerInfo
import com.jihe.punchnet.protobuf.PunchProto.SDLQueryInfo
import com.jihe.punchnet.protobuf.PunchProto.SDLRegister
import com.jihe.punchnet.protobuf.PunchProto.SDLRegisterAck
import com.jihe.punchnet.protobuf.PunchProto.SDLSendRegisterEvent
import com.jihe.punchnet.sdlan.config.EtherType
import com.jihe.punchnet.sdlan.config.SDLanConfig
import com.jihe.punchnet.sdlan.logs.TerminalLogger
import com.jihe.punchnet.sdlan.utils.isMultiBroadcast
import com.jihe.punchnet.sdlan.utils.macToString
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.async
import kotlinx.coroutines.launch
import java.nio.ByteBuffer
import java.nio.ByteOrder


suspend fun handlePacketData(
    scope: CoroutineScope,
    node: Node,
    body: ByteBuffer,
    senderSock: SDLanSock
) {
    val data: SDLData
    try {
        data = SDLData.parseFrom(body)
    } catch (e: Exception) {
        TerminalLogger.errorf { "failed to decode DATA: $e" }
        return
    }

    if (data.isP2P) {
        TerminalLogger.debugf {
            "[P2P] Rx data from ${senderSock}"
        }
        // CoroutineScope(Dispatchers.Default).async {
        scope.async {
            checkPeerRegistrationNeeded(
                node,
                false,
                data.srcMac,
                NatType.NoNat,
                senderSock,
            )
        }
    } else {
        TerminalLogger.debugf {
            "[PsP] Rx data from ${macToString(data.srcMac)} via ${senderSock}"
        }
    }
    handleTunPacket(node, senderSock, !data.isP2P, data)
}

suspend fun handleTunPacket(
    node: Node,
    senderSock: SDLanSock,
    fromSN: Boolean,
    data: SDLData,
) {
    val payload = data.data
    val origin = node.aes.decrypt(payload.toByteArray())
    if (origin == null) {
        TerminalLogger.errorf { "error to decrypt data" }
        return
    }
    TerminalLogger.debugf {"handle data from net"}
    node.iface?.handleDataFromNet(node, origin)
}

fun handlePacketRegisterACK(
    node: Node,
    body: ByteBuffer,
    senderSock: SDLanSock
) {

    val ack: SDLRegisterAck
    try {
        ack = SDLRegisterAck.parseFrom(body)
    } catch (e: Exception) {
        TerminalLogger.errorf { "failed to decode REGISTERACK: $e" }
        return
    }

    val viaMulticast = isMultiBroadcast(ack.dstMac)
    if (viaMulticast && ack.srcMac == node.mac) {
        TerminalLogger.debugf { "skip register from self" }
        return
    }

    TerminalLogger.infof {
        "Rx REGISTERACK from ${macToString(ack.srcMac)} [${senderSock}] to ${macToString(ack.dstMac)}"
    }

    peerSetP2PConfirmed(node, ack.srcMac, senderSock)
}

fun peerSetP2PConfirmed(node: Node, mac: Mac, sock: SDLanSock) {
    val scan = node.pendingPeers.remove(mac)
    if (scan == null) {
        TerminalLogger.errorf { "failed to find sender in pending peers: ${sock}" }
        return
    }
    val now = System.currentTimeMillis()/1000
    var scan2 = node.knownPeers.get(mac)
    if (scan2 == null) {
        node.knownPeers[mac] = scan
        scan2 = node.knownPeers.get(mac)
    }
    scan2?.sock = sock
    scan2?.lastP2P?.set(now)
    scan2?.lastSeen?.set(now)

    TerminalLogger.debugf {
        "P2P connection established: ${macToString(mac)} [${sock}]"
    }

    TerminalLogger.debugf {
        "===> new Peer: ${macToString(mac)} -> ${sock}"
    }

}

suspend fun handlePacketRegister(
    node: Node,
    body: ByteBuffer,
    senderSock: SDLanSock
) {
    if(!node.aes.isAuthorized()) {
        TerminalLogger.errorf { "drop register due to not authorized"}
        return
    }
    val reg: SDLRegister
    try {
        reg = SDLRegister.parseFrom(body)
    } catch (e: Exception) {
        TerminalLogger.errorf { "failed to parse REGISTER: $e"}
        return
    }
    val srcMac = reg.srcMac

    val viaMulticast = isMultiBroadcast(reg.dstMac)
    if (viaMulticast && reg.srcMac == node.mac) {
        TerminalLogger.debugf { "skip register from self" }
        return
    }

    // must be from peer

    TerminalLogger.infof {
        "[P2P] Rx REGISTER from ${senderSock}, deleting from pending"
    }
    val remoteNAT = node.pendingPeers.remove(srcMac)?.natType?:NatType.Invalid

    sendRegisterACK(node,senderSock, reg)

    checkPeerRegistrationNeeded(node, false, reg.srcMac, remoteNAT, senderSock)
}

suspend fun sendRegisterACK(
    node: Node,
    sender: SDLanSock,
    reg: SDLRegister,
) {
    if (!node.config.allowP2P) {
        TerminalLogger.debugf {
            "Skipping REGISTER ACK as P2P is not allowed"
        }
        return
    }

    // must be authed
    val ack = SDLRegisterAck.newBuilder()
        .setNetworkId(node.networkID.get())
        .setSrcMac(node.deviceConfig.mac)
        .setDstMac(reg.srcMac)
        .build()

    val data = encodeToUDPMessage(ack, PacketType.RegisterACK)
    sendToSock(node, data, sender)
}

suspend fun checkPeerRegistrationNeeded(
    node: Node,
    fromSN: Boolean,
    srcMac: Mac,
    remoteNat: NatType,
    peerSock: SDLanSock,
) {
    val peer = node.knownPeers.get(srcMac)
    if (peer == null) {
        TerminalLogger.debugf {
            "not found in known, send REGISTER to peer"
        }
        registerWithNewPeer(node,srcMac, remoteNat, peerSock, System.currentTimeMillis()/1000)
        return
    }

    val now = System.currentTimeMillis()/1000
    if (!fromSN) {
        peer.lastP2P.set(now)
    }

    if (peerSock.family != peer.sock.family) {
        TerminalLogger.errorf { "family changed, just ignore it"}
        return
    }
    if (now - peer.lastSeen.get() > 3) {
        checkKnownPeerSockChanged(node, fromSN, srcMac, peerSock, now)
    }

}
suspend fun checkKnownPeerSockChanged(
    node: Node,
    fromSN: Boolean,
    srcMac: Mac,
    peerSock: SDLanSock,
    now: Long,
) {
    if (isMultiBroadcast(srcMac)) {
        return
    }
    val peer = node.knownPeers.get(srcMac)
    if (peer == null) {
        return
    }

    if (!peerSock.is_equal(peer.sock)) {
        if (!fromSN) {
            TerminalLogger.infof {
                "peer changed: ${srcMac}: ${peer.sock} -> $peerSock"
            }
            val remoteNAT = peer.natType
            node.knownPeers.remove(srcMac)

            registerWithNewPeer(node, srcMac, remoteNAT, peerSock, now)
        }
    } else {
        peer.lastSeen.set(now)
    }
}

suspend fun registerWithNewPeer(
    node: Node,
    // fromSN: Boolean,
    mac: Mac,
    remoteNat: NatType,
    peerSock: SDLanSock,
    now: Long,
) {
    var peer = node.pendingPeers.get(mac)
    if (peer == null) {
        TerminalLogger.debugf {
            "===> new pending: ${macToString(mac)} => ${peerSock}"
        }
        node.pendingPeers[mac] = EdgePeer(0, node.deviceConfig.ip.netBitLen, peerSock, null, now)
        peer = node.pendingPeers.get(mac)
        peer?.lastSeen?.set(now)

        sendRegister(node, remoteNat, peerSock, mac)
        registerWithLocalPeers(node)
    } else {
        peer.sock = peerSock
        peer.lastSeen.set(now)
    }

}

suspend fun registerWithLocalPeers(node: Node) {
    if (!node.config.dropMulticast) {
        sendRegister(node, NatType.NoNat, node.multicastSock, SDLanConfig.BROADCASTMAC)
    }
}

suspend fun sendRegister(
    node: Node,
    natType: NatType,
    sock: SDLanSock,
    mac: Mac,
) {
    if (!node.config.allowP2P) {
        TerminalLogger.debugf { "skipping REGISTER as p2p is disabled" }
        return
    }

    if (!node.aes.isAuthorized()) {
        TerminalLogger.debugf { "skipping REGISTER as not authed" }
        return
    }

    val register = SDLRegister.newBuilder()
        .setNetworkId(node.networkID.get())
        .setSrcMac(node.deviceConfig.mac)
        .setDstMac(mac)
        .build()
    val msg = encodeToUDPMessage(register, PacketType.Register)

    sendToSock(node, msg, sock)

    // TODO: need guess port
}

suspend fun sendPacketToNet(node: Node, dstmac: Mac, content: List<Byte>, size: Long) {
    val destination = findPeerDestination(node, dstmac, size)
    TerminalLogger.debugf { "send PACKET to ${destination}" }
    sendToSock(node, content, destination)
}


suspend fun findPeerDestination(node: Node, dstmac: Mac, size: Long): SDLanSock {
    var is_p2p: Boolean = false
    var is_multicast: Boolean = false
    var result: SDLanSock
    if (isMultiBroadcast(dstmac)) {
        node.stats.txSup.addAndGet(size)
        node.stats.txBroadcast.addAndGet(size)
        result = node.config.supernode
        is_multicast = true
    } else {
        val peer = node.knownPeers.get(dstmac)
        if (peer == null) {
            node.stats.txSup.addAndGet(size)
            result = node.config.supernode
        } else {
            val now = System.currentTimeMillis()/1000
            if (now - peer.lastP2P.get() >= peer.timeout/2) {
                TerminalLogger.warning { "last p2p is too old, deleting from known hosts" }
                node.knownPeers.remove(dstmac)
                node.stats.txSup.addAndGet(size)
                result = node.config.supernode
            } else {
                is_p2p = true
                node.stats.txP2P.addAndGet(size)
                result = peer.sock
            }
        }
    }

    if(!is_p2p && !is_multicast) {
        TerminalLogger.debugf { "check_query_peer_info" }
        checkQueryPeerInfo(node, dstmac)
    }

    return result
}

suspend fun checkQueryPeerInfo(node: Node, dstmac: Mac) {
    val now = System.currentTimeMillis()/1000
    val peer = node.pendingPeers.get(dstmac)

    val needSendQuery: Boolean
    if (peer != null) {
        if (now - peer.lastSentQuery.get() > SDLanConfig.RegisterInterval) {
            needSendQuery = true
            peer.lastSentQuery.set(now)
        } else {
            needSendQuery = false
        }
    } else {
        val sock = SDLanSock(IPFamily.IPV4, 0, ByteArray(4))
        val peer = EdgePeer(
            0,
            node.deviceConfig.ip.netBitLen,
            sock,
            null,
            now,
        )
        node.pendingPeers[dstmac] = peer
        needSendQuery = true
    }

    if (needSendQuery) {
        TerminalLogger.debugf { "send query for ${macToString(dstmac)}" }
        registerWithLocalPeers(node)
        sendQueryPeer(node, dstmac)
    }

}

suspend fun sendQueryPeer(node: Node, dstmac: Mac) {
    if (!node.aes.isAuthorized()) {
        TerminalLogger.errorf { "not authed for send query" }
        return
    }
    val query = SDLQueryInfo.newBuilder()
        .setDstMac(dstmac)
        .build()
    val msg = encodeToTcpMessage(query, node.getNextPacketID(), PacketType.QueryInfo)
    node.toSocket.send(msg.toByteArray())
}

fun formEthernetPacket(srcmac: ByteArray, dstmac: ByteArray, data: ByteArray): ByteBuffer {
    val buffer = ByteBuffer.allocate(14 + data.size + 4).order(ByteOrder.BIG_ENDIAN)
    buffer.put(dstmac)
    buffer.put(srcmac)
    buffer.putShort(EtherType.IPV4)
    buffer.put(data)
    buffer.flip()
    return buffer
}

// frommac is self's mac
// fromip is self ip
// queryip is the targetip
suspend fun sendArpRequest(node: Node, queryip: Int) {
    println(1)
    val frommac = node.mac.toByteArray()
    val fromip = node.deviceConfig.ip.netAddr
    val buffer = ByteBuffer.allocate(64).order(ByteOrder.BIG_ENDIAN)
    buffer.put(ByteArray(6){0xff.toByte()})
    buffer.put(frommac)
    println(2)
    // println("src mac: ${macToString(node.mac)}")
    buffer.putShort(EtherType.ARP)
    buffer.putShort(EtherType.HWTYPE_ETH)
    buffer.putShort(EtherType.IPV4)

    buffer.put(6)
    buffer.put(4)
    buffer.putShort(EtherType.ARP_REQUEST)
    buffer.put(frommac)
    buffer.putInt(fromip)
    buffer.put(ByteArray(6))
    buffer.putInt(queryip)

    buffer.flip()
    println(3)

    val size = buffer.remaining()

    println(3.1)
    val output = node.aes.encrypt(buffer)
    println(3.2)

    println(4)
    if (output != null) {
        val d = output.toByteString()
        val data = SDLData.newBuilder()
            .setIsP2P(true)
            .setNetworkId(node.networkID.get())
            .setTtl(2)
            .setSrcMac(frommac.toByteString())
            .setDstMac(SDLanConfig.BROADCASTMAC)
            .setData(d)
            .build()
        println(5)
        val msg = encodeToUDPMessage(data, PacketType.Data)
        sendPacketToNet(node, SDLanConfig.BROADCASTMAC, msg, size.toLong())
    }
    println(6)
}

suspend fun handleTcpCommand(node: Node, cmdtype: Byte, cmdprotobuf: ByteArray) {}

suspend fun handlePacketPeerInfo(node: Node, content: ByteArray) {
    val pinfo: SDLPeerInfo
    try {
        pinfo = SDLPeerInfo.parseFrom(content)
    } catch (e: Exception) {
        TerminalLogger.errorf { "failed to decode peer info: $e"}
        return
    }

    if (pinfo.dstMac == SDLanConfig.BROADCASTMAC) {
        // pong from sn
        return
    }

    val remoteNat = NatType.fromUByte(pinfo.v4Info.natType.toUByte())
    val pending = node.pendingPeers.get(pinfo.dstMac)
    if (pending == null) {
        TerminalLogger.debugf { "Rx PEERINFO unknown peer: ${macToString(pinfo.dstMac)}" }
        return
    }
    pending.sock = SDLanSock(IPFamily.IPV4, pinfo.v4Info.port, pinfo.v4Info.v4.toByteArray())
    pending.natType = remoteNat
    TerminalLogger.debugf { "Rx PEERINFO for ${macToString(pinfo.dstMac)} is at ${pending.sock}" }
    sendRegister(node, remoteNat, pending.sock, pinfo.dstMac)
}

suspend fun handleTcpEvent(node: Node, event: EventType, cmdprotobuf: ByteArray) {
    when(event) {
        EventType.SendRegister -> {
            val reg: SDLSendRegisterEvent
            try {
                reg = SDLSendRegisterEvent.parseFrom(cmdprotobuf)
            } catch (e: Exception) {
                TerminalLogger.errorf {"failed to decode SendRegisterEvent: $e"}
                return
            }
            val remoteNat = NatType.fromUByte(reg.natType.toUByte())
            val ip = byteArrayOf(
                (reg.natIp ushr 24).and(0xff).toByte(),
                (reg.natIp ushr 16).and(0xff).toByte(),
                (reg.natIp ushr 8).and(0xff).toByte(),
                (reg.natIp).and(0xff).toByte(),
            )
            checkPeerRegistrationNeeded(node,false, reg.dstMac, remoteNat, SDLanSock(IPFamily.IPV4, reg.natPort, ip))
        }
        else -> {
            TerminalLogger.warning { "unhandled event: $event" }
        }
    }
}

fun ipInt2ByteArray(ip: Int): ByteArray {
    return byteArrayOf(
        (ip ushr 24).and(0xff).toByte(),
        (ip ushr 16).and(0xff).toByte(),
        (ip ushr 8).and(0xff).toByte(),
        ip.and(0xff).toByte(),
    )
}