package com.jihe.punchnet.sdlan.network

// import sdlanproto.Message.*
import android.util.Log
import com.google.protobuf.kotlin.toByteString
import com.jihe.punchnet.protobuf.PunchProto
import com.jihe.punchnet.protobuf.PunchProto.SDLStunProbe
import com.jihe.punchnet.protobuf.PunchProto.SDLStunProbeReply
import com.jihe.punchnet.protobuf.PunchProto.SDLStunRequest
import com.jihe.punchnet.sdlan.config.SDLanMulticastConfig
import com.jihe.punchnet.sdlan.logs.TerminalLogger
import com.jihe.punchnet.sdlan.utils.AES
import com.jihe.punchnet.sdlan.utils.RSA
import com.jihe.punchnet.sdlan.utils.UniqueNodeID
import com.jihe.punchnet.sdlan.utils.generateRandomMAC
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.SendChannel
import kotlinx.coroutines.withTimeoutOrNull
import java.net.InetSocketAddress
import java.net.SocketAddress
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicInteger
import java.util.concurrent.atomic.AtomicLong

class DeviceConfig (
    var mtu: Int,
    var mac: Mac,
    var ip: IPSubnet,
) {
    constructor(mtu: Int): this(
        mtu,
        ByteArray(6).toByteString(),
        IPSubnet(0, 0),
    )
}


class NodeConfig (
    val baseDir: String,
    val name: String,
    val nodeUUID: String,

    val allowRouting: Boolean,
    val dropMulticast: Boolean,

    val allowP2P: Boolean,
    val mtu: Int,
    val tos: Int,
    val registerSuperInterval: Byte,

    val registerTTL: Int,
    val localPort: Int,

    val supernode: SDLanSock,
    val stunServer: SDLanSock,
    val natServer1: java.net.InetSocketAddress,
    val natServer2: java.net.InetSocketAddress
)

class NodeStats private constructor(
    // communicate with p2p
    val txP2P: AtomicLong,
    val rxP2P: AtomicLong,
    // communicate with server
    val txSup: AtomicLong,
    val rxSup: AtomicLong,
    // communicate through broadcast
    val txBroadcast: AtomicLong,
    val rxBroadcast: AtomicLong,

    // last time communicate with server
    val lastSup: AtomicLong,
    // last time communicate with the edge
    val lastP2P: AtomicLong,
) {
    constructor(): this(
        AtomicLong(0),
        AtomicLong(0),
        AtomicLong(0),
        AtomicLong(0),
        AtomicLong(0),
        AtomicLong(0),
        AtomicLong(0),
        AtomicLong(0),
    )
}

class NatProbeCookie {
    val cookieid: AtomicInteger = AtomicInteger(1)
    val cookiemap = ConcurrentHashMap<Int, SendChannel<SDLStunProbeReply>>()

    fun getChannelWithID(id: Int): SendChannel<SDLStunProbeReply>? {
        return this.cookiemap.get(id)
    }

    suspend fun sendMessage(id: Int, value: SDLStunProbeReply) {
        cookiemap.remove(id)?.send(value)
    }

    fun addChannel(channel: SendChannel<SDLStunProbeReply>): Int {
        val ck = cookieid.getAndAdd(1)
        cookiemap.set(ck, channel)
        return ck
    }

    fun remove(id: Int) {
        cookiemap.remove(id)
    }
}

class Node private constructor (
    // 当前的packet id
    val packetID: AtomicInteger,
    // 自身所处的network的id
    val networkID: AtomicInteger,
    // 所处网络的token
    var token: String,

    // tun设备（vpnservice）的配置
    val deviceConfig: DeviceConfig,
    // interface的抽象
    var iface: Iface?,

    // val authorized: AtomicBoolean,

    // rsa info
    val rsa: RSA,
    // aes info
    val aes: AES,
    // var encryptKey: ByteArray,

    var config: NodeConfig,

    val pendingPeers: PeerMap,
    val knownPeers: PeerMap,

    // 自身的公网ip
    val outerIPV4: AtomicInteger,
    var udpSockV4: SDLanSocket,

    var udpSockMulticast: SDLanSocket?,
    var multicastSock: SDLanSock,

    val stats: NodeStats,

    // send message to socket, if connected, and authorized
    var toSocket: SendChannel<ByteArray>,

    var startStopChannel: SendChannel<StartStopChanInfo>,

    val mac: Mac = generateRandomMAC(),
    var nat_type: NatType = NatType.PortRestricted,
    val natProbeCookie: NatProbeCookie = NatProbeCookie(),

    var authData: com.jihe.punchnet.api.AuthResponseData? = null,
    var connectData: com.jihe.punchnet.api.ConnectResponseData? = null,
    var clientId: String = "",
    var sessionToken: com.google.protobuf.ByteString = com.google.protobuf.ByteString.EMPTY,
    val queriedPolicies: ConcurrentHashMap<Int, Int> = ConcurrentHashMap(),
    val queriedPolicyTargets: ConcurrentHashMap<Int, Int> = ConcurrentHashMap(),
    val policyRules: ConcurrentHashMap<Int, ByteArray> = ConcurrentHashMap(),
) {
    companion object {
        private var instance: Node? = null

        @Synchronized
        fun initialize(
            config: NodeConfig,
            v4Sock: SDLanSocket,
            multicastSock: SDLanSocket?,
            token: String,
            rsa: RSA,
            iface: Iface,
            toSocket: SendChannel<ByteArray>,
            startStopChannel: SendChannel<StartStopChanInfo>,
        ) : Node{
            if (instance == null) {
                instance = Node(
                    config,
                    v4Sock,
                    multicastSock,
                    token,
                    rsa,
                    iface,
                    toSocket,
                    startStopChannel,
                )
            }
            return instance!!
        }

        fun getInstance(): Node {
            return instance ?: throw IllegalStateException("Node is not initialized")
        }
    }

    private constructor(
        // pubkey: String,
        config: NodeConfig,
        v4Sock: SDLanSocket,
        multicastSock: SDLanSocket?,
        token: String,
        rsa: RSA,
        iface: Iface,
        toSocket: SendChannel<ByteArray>,
        startStopChannel: SendChannel<StartStopChanInfo>,
        // privateKey: ByteArray,
        // mtu: Int
    ): this (
        packetID = AtomicInteger(1),
        networkID = AtomicInteger(0),
        token = token,
        deviceConfig = DeviceConfig(config.mtu),
        iface = iface,
        rsa = rsa,
        aes = AES.getAES(),
        config = config,
        pendingPeers = PeerMap(),
        knownPeers = PeerMap(),

        outerIPV4 = AtomicInteger(0),
        udpSockV4 = v4Sock,
        udpSockMulticast = multicastSock,

        toSocket = toSocket,
        startStopChannel = startStopChannel,

        multicastSock = SDLanSock(
            IPFamily.IPV4,
            SDLanMulticastConfig.MULTICAST_PORT,
            SDLanMulticastConfig.MULTICAST_V4,
        ),

        stats = NodeStats(),
    ) {
        this.deviceConfig.mac = this.mac
    }

    fun getNextPacketID(): Int {
        return packetID.getAndAdd(1)
    }

    suspend fun sendStunRequest() {
        if (this.clientId.isEmpty()) {
            this.clientId = "punchnet_android_client_" + UniqueNodeID.getUUID()
        }
        val req = SDLStunRequest.newBuilder()
            .setClientId(this.clientId)
            .setNetworkId(this.networkID.get())
            .setIp(this.deviceConfig.ip.netAddr)
            .setMac(this.mac)
            .setNatType(this.nat_type.toByte().toInt())
            .setV6Info(PunchProto.SDLV6Info.getDefaultInstance())
            .setSessionToken(this.sessionToken)
            .build()

        val msg = encodeToUDPMessage(req, PacketType.StunRequest)
        TerminalLogger.debugf {
            "send STUN_REQUEST to ${config.stunServer}, clientIdTail=${clientId.takeLast(12)}, " +
                "network=${req.networkId}, ip=${com.jihe.punchnet.sdlan.utils.ipToString(req.ip)}, " +
                "mac=${com.jihe.punchnet.sdlan.utils.macToString(req.mac)}, natType=${nat_type}, " +
                "sessionTokenBytes=${sessionToken.size()}"
        }
        sendToSock(this, msg, config.stunServer)
    }



    suspend fun ping_to_sn() {
        val msg = encodeToControlMessage(null, PacketType.Ping).toByteArray()
        _sendDataToSocket(msg)
    }

    suspend fun _sendDataToSocket(msg: ByteArray) {

        if (aes.isAuthorized()) {
            TerminalLogger.debugf{"authorized, ping to sn"}
            toSocket.send(msg)
        } else {
            TerminalLogger.debugf{"unauthorized, not ping to sn"}
        }
    }

    suspend fun probeNatType() {
        val previousNatType = nat_type
        val reply1 = this._sendAndWaitForProbeReply(StunProbeAttr.None, config.natServer1)
        if (reply1 == null) {
            TerminalLogger.warning { "STUN probe attr=None timed out; keeping natType=${nat_type}" }
            return
        }
        if (reply1.ip == (outerIPV4.get())) {
           if (this._sendAndWaitForProbeReply(StunProbeAttr.Peer, config.natServer1) == null) {
               // failed to get with peer, just symmetric
               nat_type = NatType.Symmetric
               return
           }
            nat_type = NatType.NoNat
            return
        }

        val reply2 = this._sendAndWaitForProbeReply(StunProbeAttr.Peer, config.natServer1)
        if (reply2 != null) {
            nat_type = NatType.FullCone
            return
        }

        val reply3 = this._sendAndWaitForProbeReply(StunProbeAttr.None, config.natServer2)
        if (reply3 == null) {
            nat_type = if (previousNatType == NatType.Invalid || previousNatType == NatType.Blocked) {
                NatType.PortRestricted
            } else {
                previousNatType
            }
            TerminalLogger.warning {
                "secondary STUN probe timed out after primary reply; keeping natType=${nat_type}"
            }
            return
        }

        if ((reply1.ip != reply3.ip) || (reply1.port != reply3.port)) {
            nat_type = NatType.Symmetric
            return
        }

        val reply4 = this._sendAndWaitForProbeReply(StunProbeAttr.Port, config.natServer1)
        if (reply4 == null) {
            nat_type = NatType.PortRestricted
        } else {
            nat_type = NatType.ConeRestrict
        }

    }

    suspend fun _sendAndWaitForProbeReply(attr: StunProbeAttr, toServer: java.net.SocketAddress): SDLStunProbeReply? {
        val channel = kotlinx.coroutines.channels.Channel<SDLStunProbeReply>(100)
        val cookie = natProbeCookie.addChannel(channel)
        val probe = com.jihe.punchnet.protobuf.PunchProto.SDLStunProbe.newBuilder()
            .setAttr(attr.ordinal.toInt())
            .setCookie(cookie)
            .build()

        val msg = encodeToUDPMessage(probe, PacketType.StunProbe)
        TerminalLogger.debugf { "send STUN_PROBE attr=$attr cookie=$cookie to $toServer" }
        this.udpSockV4.send_to(msg.toByteArray(), toServer)

        val k = kotlinx.coroutines.withTimeoutOrNull(5000) {
            val response = channel.receive()
            return@withTimeoutOrNull response
        }
        if (k == null) {
            TerminalLogger.warning { "STUN_PROBE attr=$attr cookie=$cookie to $toServer timed out" }
        } else {
            TerminalLogger.debugf { "STUN_PROBE attr=$attr cookie=$cookie reply ip=${k.ip} port=${k.port}" }
        }
        natProbeCookie.remove(cookie)
        return k
    }
}
