package com.jihe.punchnet.sdlan.network

import android.util.Log
import com.jihe.punchnet.PunchnetServiceArgument
import com.jihe.punchnet.data.ButtonRepository
import com.jihe.punchnet.data.ButtonState

import com.jihe.punchnet.protobuf.PunchProto.SDLRegisterSuper
import com.jihe.punchnet.protobuf.PunchProto.SDLRegisterSuperAck
import com.jihe.punchnet.protobuf.PunchProto.SDLRegisterSuperNak
import com.jihe.punchnet.protobuf.PunchProto.SDLExposedServiceResponse
import com.jihe.punchnet.protobuf.PunchProto.SDLPolicyResponse
import com.jihe.punchnet.protobuf.PunchProto.SDLStunProbeReply
import com.jihe.punchnet.sdlan.config.Arguments
import com.jihe.punchnet.sdlan.config.RSAConfig
import com.jihe.punchnet.sdlan.config.SDLanConfig
import com.jihe.punchnet.sdlan.config.SDLanMulticastConfig
import com.jihe.punchnet.sdlan.config.toIPV4String
import com.jihe.punchnet.sdlan.logs.TerminalLogger
import com.jihe.punchnet.sdlan.utils.RSA
import com.jihe.punchnet.sdlan.utils.UniqueNodeID
import com.jihe.punchnet.sdlan.utils.ipToString
import com.jihe.punchnet.sdlan.utils.macToString
import com.jihe.punchnet.sdlan.utils.parseScoketAddressV4FromString
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.NonCancellable
import kotlinx.coroutines.async
import kotlinx.coroutines.cancel
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.ReceiveChannel
import kotlinx.coroutines.channels.SendChannel
import kotlinx.coroutines.delay
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import kotlinx.coroutines.selects.select
import kotlinx.coroutines.withContext
import java.io.DataInputStream
import java.io.DataOutputStream
import java.net.InetSocketAddress
import java.net.Socket
import java.net.SocketAddress
import java.nio.ByteBuffer
import java.nio.file.Paths
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicLong
import kotlin.io.path.pathString
import kotlin.system.exitProcess

suspend fun onMessage(scope: CoroutineScope, data: SDLanControl) {
    val node = Node.getInstance()

    TerminalLogger.debugf {"message received"}
    when(data.packetType) {
        PacketType.Welcome -> {
            TerminalLogger.debugf { "got welcome, sending register super" }
            if (node.clientId.isEmpty()) {
                node.clientId = "punchnet_android_client_" + UniqueNodeID.getUUID()
            }
            val registerIp = com.jihe.punchnet.sdlan.utils.ipStringToInt(node.connectData?.ip ?: "0.0.0.0").toUInt().toInt()
            
            val registerSuper = SDLRegisterSuper.newBuilder()
                .setClientId(node.clientId)
                .setNetworkId(node.authData?.networkId ?: 0)
                .setMac(com.google.protobuf.ByteString.copyFrom(node.mac.toByteArray()))
                .setIp(registerIp)
                .setMaskLen(node.connectData?.maskLen ?: 0)
                .setHostname(node.connectData?.hostname ?: "")
                .setPubKey(node.rsa.getPublicKeyString())
                .setAccessToken(node.authData?.accessToken ?: "")
                .build()

            val msg = encodeToControlMessage(registerSuper, PacketType.RegisterSuper)
            node.toSocket.send(msg.toByteArray())
            TerminalLogger.debugf {
                "sent RegisterSuper clientIdTail=${node.clientId.takeLast(12)}, network=${registerSuper.networkId}, " +
                    "ip=${ipToString(registerSuper.ip)}, mask=${registerSuper.maskLen}, mac=${macToString(registerSuper.mac)}, " +
                    "identity=${node.connectData?.identityId ?: 0}"
            }
        }
        PacketType.RegisterSuperACK -> {
            val ack = SDLRegisterSuperAck.parseFrom(data.currentPacket)
            val algorithm = ack.algorithm.ifEmpty { "aes" }.lowercase()
            TerminalLogger.debugf {
                "got register super ack: algorithm=$algorithm, region=${ack.regionId}, sessionTokenBytes=${ack.sessionToken.size()}"
            }
            if (algorithm != "aes" && algorithm != "chacha20") {
                TerminalLogger.errorf { "unsupported encryption algorithm from server: ${ack.algorithm}" }
                node.startStopChannel.send(StartStopChanInfo(StartStopFlag.IsStop, null))
                ButtonRepository.updateState(ButtonState.ButtonStopped)
                scope.cancel()
                return
            }

            val aeskey = node.rsa.decrypt(ack.key.toByteArray())
            if (aeskey == null) {
                println("aes key error")
                return
            }
            node.aes.setSecret(aeskey, algorithm, ack.regionId.toLong())
            node.sessionToken = ack.sessionToken
            
            node.deviceConfig.ip.netAddr = com.jihe.punchnet.sdlan.utils.ipStringToInt(node.connectData?.ip ?: "0.0.0.0").toUInt().toInt()
            node.deviceConfig.ip.netBitLen = (node.connectData?.maskLen ?: 24).toByte()

            if (node.iface is IfaceTun) {
                val tun = node.iface as IfaceTun
                val maskVal = com.jihe.punchnet.sdlan.network.maskDigitToInt(node.deviceConfig.ip.netBitLen.toInt()) ?: 0xffffff00.toInt()
                tun.arpTable.routeTable.addRoute(
                    com.jihe.punchnet.sdlan.network.RouteDetail(
                        mask = maskVal,
                        gw = 0,
                        maskedAddr = node.deviceConfig.ip.netAddr and maskVal
                    )
                )
            }

            node.iface?.reload_config(node.deviceConfig)
            node.networkID.set(node.authData?.networkId ?: 0)
            TerminalLogger.debugf {
                "control registered: clientIdTail=${node.clientId.takeLast(12)}, network=${node.networkID.get()}, " +
                    "ip=${ipToString(node.deviceConfig.ip.netAddr)}, mask=${node.deviceConfig.ip.netBitLen}, " +
                    "mac=${macToString(node.mac)}, identity=${node.connectData?.identityId ?: 0}"
            }
            sendExposedServiceRequest(node)
            node.sendStunRequest()
            scope.launch {
                try {
                    node.probeNatType()
                    node.sendStunRequest()
                    TerminalLogger.debugf { "updated NAT type after probe: ${node.nat_type}" }
                } catch (e: Exception) {
                    TerminalLogger.warning { "NAT probe failed: $e" }
                }
            }
        }
        PacketType.RegisterSuperNAK -> {
            val nak = SDLRegisterSuperNak.parseFrom(data.currentPacket)
            TerminalLogger.debugf { "got register super nak: $nak" }
            val nakcode = NakMsgCode.fromByte(nak.errorCode.toByte())
            when(nakcode) {
                NakMsgCode.InvalidToken, NakMsgCode.NodeDisabled-> {
                    node.startStopChannel.send(StartStopChanInfo(StartStopFlag.IsStop, null))
                    ButtonRepository.updateState(ButtonState.ButtonStopped)
                    scope.cancel()
                }
                else -> {
                    ButtonRepository.updateState(ButtonState.ButtonStopped)
                    node.startStopChannel.send(StartStopChanInfo(StartStopFlag.IsStop, null))
                }
            }
            node.aes.setSecret(null)
        }
        PacketType.Command -> {
            TerminalLogger.debugf {"Command received"}
            if (data.currentPacket.size < 1) {
                TerminalLogger.errorf { "malformed command received"}
                return
            }
            // handleTcpCommand(node, data.currentPacket[0], data.currentPacket.sliceArray(1..data.currentPacket.size-1))
        }
        PacketType.PeerInfo -> {
            TerminalLogger.debugf { "PeerInfo received" }
            handlePacketPeerInfo(node, data.currentPacket)
        }
        PacketType.Event -> {
            TerminalLogger.debugf { "Event received" }
            handleEvent(node, data.currentPacket)
        }
        PacketType.Pong -> {
            TerminalLogger.debugf { "Pong Received" }
        }
        PacketType.ArpResponse -> {
            TerminalLogger.debugf { "ArpResponse received" }
            handlePacketArpResponse(node, data.currentPacket)
        }
        PacketType.PolicyResponse -> {
            val policy = SDLPolicyResponse.parseFrom(data.currentPacket)
            node.queriedPolicies[policy.dstIdentityId] = policy.version
            node.policyRules[policy.dstIdentityId] = policy.rules.toByteArray()
            val ruleSummary = describePolicyRules(policy.rules.toByteArray())
            val targetIp = node.queriedPolicyTargets[policy.dstIdentityId]
            val targetText = targetIp?.let { ", dstIp=${ipToString(it)}, ${describeResourcesForTarget(node, it)}" } ?: ""
            if (policy.rules.isEmpty) {
                TerminalLogger.warning {
                    "PolicyResponse has no allow rules: srcIdentity=${policy.srcIdentityId}, " +
                        "dstIdentity=${policy.dstIdentityId}, version=${policy.version}$targetText"
                }
            } else {
                TerminalLogger.debugf {
                    "PolicyResponse received srcIdentity=${policy.srcIdentityId}, dstIdentity=${policy.dstIdentityId}, " +
                        "version=${policy.version}, rules=$ruleSummary$targetText"
                }
            }
        }
        PacketType.ExposedServiceResponse -> {
            val exposed = SDLExposedServiceResponse.parseFrom(data.currentPacket)
            TerminalLogger.debugf {
                "ExposedServiceResponse received version=${exposed.version}, " +
                    "tcpPorts=${exposed.tcpPortsList.size}, udpPorts=${exposed.udpPortsList.size}"
            }
        }
        else -> {
            println("error packet type: ${data.packetType.toUByte()}")
        }
    }
}

suspend fun run_sdlan(scope: CoroutineScope, iface: Iface, argument: Arguments, authData: com.jihe.punchnet.api.AuthResponseData, connectData: com.jihe.punchnet.api.ConnectResponseData) {
    UniqueNodeID.setBaseDir(argument.baseDir)
    val edgeUUID = UniqueNodeID.getUUID()
    val config = parseConfig(edgeUUID, argument)
    if (config == null) {
        TerminalLogger.errorf {
            "parse config failed"
        }
        return
    }

    val toSocket = Channel<ByteArray>(100)
    val start_stop_channel = Channel<StartStopChanInfo>(100)
    initEdge(scope, iface, argument.token, config, toSocket, start_stop_channel)

    val snParts = argument.sn.split(":")

    val now = System.currentTimeMillis() / 1000

    val node = Node.getInstance()
    node.authData = authData
    node.connectData = connectData
    TerminalLogger.debugf { "self mac: ${macToString(node.mac)}"}
    TerminalLogger.debugf { describeConnectDataForLog(node) }

    val onConnected: suspend (tech.kwik.core.QuicStream) -> Unit = { stream ->
        TerminalLogger.debugf { "QUIC stream connected callback" }
        // We wait for server's PACKET_WELCOME before sending RegisterSuper
    }

    scope.launch {
        try {
            while(true) {
                if (node.iface == null) {
                    delay(1000)
                    continue
                }
                val data = node.iface?.recv()
                Log.d("SDLAN", "async receive data from iface: ${data?.size} bytes")
                if (data == null) {
                    delay(1000)
                    continue
                }
                if (data.isEmpty()) {
                    delay(1000)
                    continue
                }
                Log.d("SDLAN", "handle data form device starts")
                node.iface?.handleDataFromDevice(node, data)
                Log.d("SDLAN", "handle data form device stops")
            }
        } catch (e: Exception) {
            println("iface read is cancelled")
            node.iface?.close()
            node.udpSockV4.close()
            node.udpSockMulticast?.close()
        }
    }

    scope.launch {
        initQUICConn(
            scope,
            snParts[0], snParts[1].toIntOrNull() ?: 1265,
            start_stop_channel,
            AtomicLong(now),
            AtomicBoolean(false),
            toSocket as ReceiveChannel<ByteArray>,
            onConnected,
            ::onMessage,
            suspend {
                node.aes.setSecret(null)
            },
            null,
        )
    }

    println("sending start")
    start_stop_channel.send(StartStopChanInfo(StartStopFlag.IsStart, null))
    println("sent start")

    val cancel = Channel<Boolean>(100)
    runEdgeLoop(scope, node, cancel)

    while(true) {
        TerminalLogger.debugf {"ping to sn"}
        delay(SDLanConfig.TCP_PING_TIME*1000)
        node.ping_to_sn()
    }
}

suspend fun runEdgeLoop(scope: CoroutineScope, node: Node, cancel: ReceiveChannel<Boolean>) {
    node.ping_to_sn()
    scope.launch {
        loopSocketV4(scope, node, node.udpSockV4, cancel)
    }
}

suspend fun loopSocketV4(scope: CoroutineScope, node: Node, sock: SDLanSocket, cancel: ReceiveChannel<Boolean>) {
    val job_stun_request = scope.async {
        try {
            while(true) {
                delay(10_000)
                node.probeNatType()
                node.sendStunRequest()
            }
        } catch (e: Exception) {
            Log.e("Socket", "job stun request exits: $e")
        }
    }

    val job_handle_packet = scope.async {
        try {
            while(true) {
                readAndParsePacket(scope, node, sock)
            }
        } catch (e: Exception) {
            sock.close()
        }
    }

    select<Unit> {
        job_stun_request.onAwait() {

        }
        cancel.onReceive() {

        }
        job_handle_packet.onAwait() {

        }
    }
    job_stun_request.cancelAndJoin()
    job_handle_packet.cancelAndJoin()
    Log.e("Socket", "loop socket v4 stops")
}

suspend fun readAndParsePacket(scope: CoroutineScope, node: Node, sock: SDLanSocket) {
    val packet = sock.receive()
    val from = packet.socketAddress

    if (packet.length < 1) {
        TerminalLogger.errorf { "got zero-sized packet"}
        return
    }
    val data = ByteBuffer.wrap(packet.data, 0, packet.length)
    // val data = packet.data.toByteString(0, packet.length)
    handleAPacket(scope, node, from, data)
}

suspend fun handleAPacket(scope: CoroutineScope, node: Node, from: SocketAddress, data: ByteBuffer) {
    val pktType = PacketType.fromValue(data.get().toUByte())
    if (pktType == null) {
        TerminalLogger.errorf { "invalid packet type" }
        return
    }
    // val buffer = ByteBuffer.wrap(data, 1, size-1)

            when (pktType) {
        PacketType.Data -> {
            TerminalLogger.debugf { "got DATA" }
            if (!node.aes.isAuthorized()) {
                TerminalLogger.errorf { "drop DATA before authed" }
                return
            }
            if (from is InetSocketAddress) {
                TerminalLogger.debugf {"got data"}
                val sock = SDLanSock(IPFamily.IPV4, from.port, from.address.address)
                handlePacketData(scope, node, data, sock)
            }
        }
        PacketType.StunProbeReply -> {
            try {
                val reply = SDLStunProbeReply.parseFrom(data)
                node.natProbeCookie.sendMessage(reply.cookie, reply)
            } catch (e: Exception) {
                TerminalLogger.errorf { "failed to decode Probe Reply"}
            }
        }
        PacketType.StunReply -> {
            TerminalLogger.debugf { "got stun reply" }
            return
        }
        PacketType.PolicyResponse -> {
            TerminalLogger.debugf { "got policy response" }
            return
        }
        PacketType.ExposedServiceResponse -> {
            TerminalLogger.debugf { "got exposed service response" }
            return
        }
        PacketType.Register -> {
            TerminalLogger.debugf { "got REGISTER" }
            if (from is InetSocketAddress) {
                val sock = SDLanSock(IPFamily.IPV4, from.port, from.address.address)
                handlePacketRegister(node, data, sock)
            }
        }
        PacketType.RegisterACK -> {
            TerminalLogger.debugf { "got REGISTERACK" }
            if (!node.aes.isAuthorized()) {
                TerminalLogger.errorf { "drop REGISTERACK before authed" }
                return
            }
            if (from is InetSocketAddress) {
                val sock = SDLanSock(IPFamily.IPV4, from.port, from.address.address)
                handlePacketRegisterACK(node, data, sock)
            }

        }
        else -> {
            TerminalLogger.debugf {"ignore packet type: $pktType"}
        }
    }
}


fun initEdge(scope: CoroutineScope, iface: Iface, token: String, config: NodeConfig, toSocket: SendChannel<ByteArray>, startStopChannel: SendChannel<StartStopChanInfo>) {
    val rsa = RSA.getRSA()
    val pathname = Paths.get(config.baseDir, RSAConfig.BASE_DIR).pathString
    Log.d("DIR", "pathname = $pathname")
    rsa.generateKeyPair(pathname)
    //rsa.generateKeyPair(Path.of(config.baseDir, RSAConfig.BASE_DIR).name)

    val sockV4 = SDLanSocket(scope, "0.0.0.0", config.localPort)

    var sockMulticast: SDLanSocket? = null
    if (!config.dropMulticast) {
        sockMulticast = SDLanSocket(scope, SDLanMulticastConfig.MULTICAST_V4.toIPV4String(), SDLanMulticastConfig.MULTICAST_PORT)
    }

    Node.initialize(
        config,
        sockV4,
        sockMulticast,
        token,
        rsa,
        iface,
        toSocket,
        startStopChannel,
    )
    val instance = Node.getInstance()
    instance.token = token
    instance.udpSockV4 = sockV4
    instance.udpSockMulticast = sockMulticast
    instance.config = config
    instance.toSocket = toSocket
    instance.startStopChannel = startStopChannel
    println(instance)
}


fun parseConfig(nodeuuid: String, argument: Arguments): NodeConfig? {
    if (argument.sn.length == 0) {
        println("no sn is specified")
        return null
    }


    val sn = parseScoketAddressV4FromString(argument.sn)
    if (sn == null) {
        return null
    }

    val stun = parseScoketAddressV4FromString(argument.stun)
    if (stun == null) {
        return null
    }

    val natServer1 = parseScoketAddressV4FromString(argument.nat_server1)
    val natServer2 = parseScoketAddressV4FromString(argument.nat_server2)
    if (natServer1 == null || natServer2 == null) {
        return null
    }

    return NodeConfig(
        baseDir = argument.baseDir,
        name = argument.name,
        nodeUUID = nodeuuid,
        allowRouting = argument.allowRouting,
        dropMulticast = argument.dropMulticast,
        allowP2P = argument.allowP2P,
        mtu = argument.mtu,
        tos = argument.tos,
        registerSuperInterval = SDLanConfig.RegisterSuperInterval,
        registerTTL = argument.registerTTL,
        localPort = argument.localPort,
        supernode = SDLanSock(IPFamily.IPV4, sn.port, sn.address.address),
        stunServer = SDLanSock(IPFamily.IPV4, stun.port, stun.address.address),
        natServer1 = natServer1,
        natServer2 = natServer2
    )
}
