package com.jihe.punchnet.sdlan.network

import android.util.Log
import com.jihe.punchnet.PunchnetServiceArgument
import com.jihe.punchnet.protobuf.PunchProto.SDLDevAddr
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.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.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.async
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: SDLanTCP) {
    val node = Node.getInstance()

    TerminalLogger.debugf {"message received"}
    when(data.packetType) {
        PacketType.RegisterSuperACK -> {
            val ack = SDLRegisterSuperAck.parseFrom(data.currentPacket)
            // TerminalLogger.debugf { "got register super ack: $ack" }
            println("ack is ${ack}")
            println("ack's key is ${ack.aesKey.size()}")

            val aeskey = node.rsa.decrypt(ack.aesKey.toByteArray())
            if (aeskey == null) {
                println("aes key error")
                return
            }
            println("decrypted aes is: ${aeskey.contentToString()}")
            node.aes.setSecret(aeskey)
            val ipStr = ipToString(ack.devAddr.netAddr)
            TerminalLogger.debugf { "assigned ip: $ipStr" }
            node.deviceConfig.ip.netAddr = ack.devAddr.netAddr
            node.deviceConfig.ip.netBitLen = ack.devAddr.netBitLen.toByte()

            println("reloading config: node.iface = ${node.iface}")
            node.iface?.reload_config(node.deviceConfig)
            println("reloading config ok: node.iface = ${node.iface}")

            node.networkID.set(ack.devAddr.networkId)
            // println("got aes key: ${aeskey.toList()}, length is ${aeskey.size}")
            node.sendStunRequest()

            scope.async {
                node.probeNatType()
                TerminalLogger.debugf { "nat type is ${node.nat_type}"}
            }
            /*
            CoroutineScope(Dispatchers.Default).async {
                node.probeNatType()
                TerminalLogger.debugf { "nat type is ${node.nat_type}"}
            }
             */
        }
        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))
                    exitProcess(-1)
                }
                else -> {
                    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" }
            if (data.currentPacket.size < 1) {
                TerminalLogger.errorf { "malformed event received"}
                return
            }
            val eventType = EventType.fromByte(data.currentPacket[0])

            val size = data.currentPacket.size
            val content = data.currentPacket.sliceArray(1..size-1)
            handleTcpEvent(node, eventType, content)
        }
        PacketType.Pong -> {
            TerminalLogger.debugf { "tcp Pong Received" }
            // node.ping_to_sn()
        }
        else -> {
            println("error packet type: ${data.packetType.toUByte()}")
        }
    }

}

suspend fun run_sdlan(scope: CoroutineScope, start_stop_channel: Channel<StartStopChanInfo>, iface: Iface, argument: Arguments, routeinfo: PunchnetServiceArgument?) {
    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 tcp = argument.tcp.split(":")

    val now = System.currentTimeMillis() / 1000

    val node = Node.getInstance()


    val before = ubyteArrayOf(126u,162u,25u,63u,148u,147u,198u,41u,69u,165u,149u,101u,153u,82u,190u,21u,48u,120u,26u,64u,142u,103u,159u,60u,47u,129u,176u,17u,232u,210u,36u,56u)
    val encrypted = node.rsa.encrypt(before.toByteArray(), use_private_key = false)
    // println("encrypted: ${encrypted.contentToString()}")
    val origin = node.rsa.decrypt(encrypted, use_pub_key = false)
    // TerminalLogger.debugf { "got encrypted ${encrypted.size}"}
    // TerminalLogger.debugf { "got origin ${origin.contentToString()}"}

    TerminalLogger.debugf { "self mac: ${macToString(node.mac)}"}


    val onConnected: suspend (Socket, Int?) -> Unit = { socket, pktID ->
        println("connected callback")
        val clientid = UniqueNodeID.getUUID()
        println("client id is $clientid")

        val registerSuper = SDLRegisterSuper.newBuilder()
            .setVersion(1)
            .setToken(node.token)
            .setInstalledChannel("kotlin")
            .setClientId(clientid)
            .setDevAddr(
                SDLDevAddr.newBuilder().setMac(node.mac)
                    .setNetAddr(0)
                    .setNetBitLen(0)
                    .setNetworkId(0)
                    .build()
            )
            .setPubKey(node.rsa.getPublicKeyString())
            .build()

        val packetID = pktID ?: node.getNextPacketID()

        val msg = encodeToTcpMessage(registerSuper, packetID, PacketType.RegisterSuper)
        // TerminalLogger.debugf{"sent register super: $registerSuper"}
        toSocket.send(msg.toByteArray())
        TerminalLogger.debugf{"sent register super!"}
    }
    /*
    val onMessage: suspend (SDLanTCP) -> Unit = { data ->
    }
     */

    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)
                    delay(1000)
                    // println("got data is null")
                    continue
                }
                if (data.isEmpty()) {
                    delay(1000)
                    // println("got data size 0")
                    // 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()
        }
    }
    /*
    CoroutineScope((Dispatchers.IO)).async {
        while(true) {
            val data = node.iface?.recv()
            Log.d("SDLAN", "async receive data from iface: ${data?.size} bytes")
            if (data == null) {
                //delay(1000)
                println("got data is null")
                continue
            }
            if (data.isEmpty()) {
                println("got data size 0")
                // delay(1000)
                continue
            }
            Log.d("SDLAN", "handle data form device starts")
            node.iface?.handleDataFromDevice(node, data)
            Log.d("SDLAN", "handle data form device stops")
        }
    }
     */

    // CoroutineScope(Dispatchers.IO).async {
    scope.launch {
        try {
            initTCPConn(
                scope,
                tcp[0], tcp[1].toInt(),
                start_stop_channel,
                AtomicLong(now),
                AtomicBoolean(false),
                toSocket as ReceiveChannel<ByteArray>,
                onConnected,
                ::onMessage,
                suspend {
                    node.aes.setSecret(null)
                },
                null,
            )
        } finally {
            println("initTCPConn is cancelled")
        }
    }


    start_stop_channel.send(StartStopChanInfo(StartStopFlag.IsStart, null))

    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 {
        while(true) {
            delay(10_000)
            node.sendStunRequest()
        }
    }

    val job_handle_packet = scope.async {
        while(true) {
            readAndParsePacket(scope, node, sock)
        }
    }

    select<Unit> {
        job_stun_request.onAwait() {

        }
        cancel.onReceive() {

        }
        job_handle_packet.onAwait() {

        }
    }
    job_stun_request.cancelAndJoin()
    job_handle_packet.cancelAndJoin()
}

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.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()
    println(instance)
}

suspend fun initTCPConn(
    scope: CoroutineScope,
    tcpHost: String,
    tcpPort: Int,
    start_stop: Channel<StartStopChanInfo>,
    pong_time: AtomicLong,
    connected: AtomicBoolean,
    toSocket: ReceiveChannel<ByteArray>,
    onConnected: suspend (stream: Socket, pktID: Int?)->Unit,
    onMessage: suspend (CoroutineScope, SDLanTCP)->Unit,
    onDisconnected: suspend ()->Unit,
    connectingChan: SendChannel<ConnectingState>?
) {
    val started = AtomicBoolean(false)
    // var started: Boolean = false
    var startPktID: Int? = null

    while(true) {
        connectingChan?.send(ConnectingState.NotConnected)

        if (!started.get()) {
            while(true) {
                println("waiting for start_stop")
                val startStopInfo = start_stop.receive()
                if (startStopInfo.flag == StartStopFlag.IsStart) {
                    started.set(true)
                    // started = true
                    startPktID = startStopInfo.packetID
                    break
                }
                TerminalLogger.debugf {
                    "start stop chan received ${startStopInfo}"
                }
            }
        }

        connectingChan?.send(ConnectingState.Connecting)
        TerminalLogger.debugf {"try connecting tcp..."}

        val socket: Socket
        try {
            withContext(Dispatchers.IO) {
                TerminalLogger.debugf { "connecting to $tcpHost:$tcpPort" }
                socket = Socket(tcpHost, tcpPort)
            }
        } catch (e: Exception) {
            TerminalLogger.errorf { "failed to connect to ${tcpHost}:${tcpPort}: $e"}
            delay(3000)
            continue
        }

        val node = Node.getInstance()
        val outIP = ByteBuffer.wrap(socket.localAddress.address).getInt()
        node.outerIPV4.set(outIP)

        // val job_read_packet = CoroutineScope(Dispatchers.IO).async {
        val job_read_packet = scope.async {
            val input = DataInputStream(socket.getInputStream())
            try {
                println("job read packet starts")
                while(true) {
                    val tcpPacket = readPacket(input)
                    if (tcpPacket == null) {
                        TerminalLogger.errorf {"tcp Packet is null"}
                        break
                    }
                    onMessage(scope, tcpPacket)
                }
            } finally {
                TerminalLogger.errorf {"input closing"}
                input.close()
            }
        }

        // val job_write_to_packet = CoroutineScope(Dispatchers.IO).async {
        val job_write_to_packet = scope.async {
            val output = DataOutputStream(socket.getOutputStream())
            try {
                TerminalLogger.debugf {"job write to packet starts"}
                while(true) {
                    try {
                        val msg = toSocket.receive()
                        TerminalLogger.debugf{"received message"}
                        output.write(msg)
                    } catch (e: Exception) {
                        output.close()
                        TerminalLogger.errorf {"failed to receive message to tcp: $e"}
                        break
                    }
                }
            } finally {
                output.close()
            }
        }

        // val job_check_pong = CoroutineScope(Dispatchers.IO).async {
        val job_check_pong = scope.async {
            println("job check pong starts")
            while(true) {
                delay(10_000)
                val now = System.currentTimeMillis()/1000
                if (connected.get() && (now-pong_time.get()>SDLanConfig.TCP_PING_TIME*2)) {
                    TerminalLogger.errorf {"tcp pong check expired"}
                    break
                }
            }
        }

        // val job_check_stop = CoroutineScope(Dispatchers.IO).async {
        val job_check_stop = scope.async {
            println("job check stop starts")
            while(true) {
                try {
                    val v = start_stop.receive()
                    if (v.flag == StartStopFlag.IsStop) {
                        started.set(false)
                        break
                    }
                } catch(e: Exception) {
                    started.set(false)
                    break
                }

            }
        }

        TerminalLogger.debugf { "connected" }
        onConnected(socket, startPktID)
        connectingChan?.send(ConnectingState.Connected)

        var cancelled: Boolean = false
        select<Unit> {
            job_read_packet.onAwait() {
                println("job read packet exited")
            }
            job_write_to_packet.onAwait() {
                println("job write to packet exited")
            }
            job_check_pong.onAwait() {
                println("job check pong exited")
            }
            job_check_stop.onAwait() {
                println("job check stop exited")
            }
        }

        println("m1")
        job_read_packet.cancelAndJoin()
        println("m2")
        job_write_to_packet.cancelAndJoin()
        println("m3")
        job_check_pong.cancelAndJoin()
        println("m4")
        job_check_stop.cancelAndJoin()
        println("m5")

        delay(1000)
    }
}

fun readPacket(input: DataInputStream): SDLanTCP? {
    try {
        val size = input.readShort()
        val packetID = input.readInt().toUInt()
        val packetTypeUByte = input.readByte().toUByte()
        val packetType = PacketType.fromValue(packetTypeUByte)
        if (packetType == null) {
            TerminalLogger.errorf{"packet type not found: $packetTypeUByte"}
            return null
        }

        if (size < 5) {
            TerminalLogger.errorf {"got input stream size error: ${size}"}
            return null
        }

        var buffsize = size - 5
        val data = ByteArray(buffsize)
        var toread = buffsize
        while(toread > 0) {
            val sizeGot = input.read(data, (buffsize-toread), toread)
            if (sizeGot <= 0) {
                TerminalLogger.errorf {"failed to read further: got ${sizeGot}"}
                return null
            }
            toread -= sizeGot
        }
        return SDLanTCP(packetID, packetType, data)
    } catch (e: Exception) {
        TerminalLogger.errorf { "failed to read: $e" }
        return null
    }
}

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

    val natServer1 = parseScoketAddressV4FromString(argument.nat_server1)
    val natServer2 = parseScoketAddressV4FromString(argument.nat_server2)

    if (natServer1 == null || natServer2 == null) {
        return null
    }

    val sn = parseScoketAddressV4FromString(argument.sn)
    if (sn == 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),
        natServer1 = natServer1,
        natServer2 = natServer2,
    )
}