package com.jihe.punchnet.sdlan.network

import android.util.Log
import com.jihe.punchnet.sdlan.logs.TerminalLogger
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
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 tech.kwik.core.QuicClientConnection
import tech.kwik.core.QuicStream
import java.net.URI
import java.nio.ByteBuffer
import java.io.InputStream
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicLong

suspend fun initQUICConn(
    scope: CoroutineScope,
    quicHost: String,
    quicPort: Int,
    start_stop: Channel<StartStopChanInfo>,
    pong_time: AtomicLong,
    connected: AtomicBoolean,
    toSocket: ReceiveChannel<ByteArray>,
    onConnected: suspend (stream: QuicStream)->Unit,
    onMessage: suspend (CoroutineScope, SDLanControl)->Unit,
    onDisconnected: suspend ()->Unit,
    connectingChan: SendChannel<ConnectingState>?
) {
    val started = AtomicBoolean(false)
    val node = Node.getInstance()

    while (scope.isActive) {
        connectingChan?.send(ConnectingState.NotConnected)

        if (!started.get()) {
            while (true) {
                val startStopInfo = start_stop.receive()
                if (startStopInfo.flag == StartStopFlag.IsStart) {
                    started.set(true)
                    break
                }
            }
        }

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

        var connection: QuicClientConnection? = null
        var stream: QuicStream? = null

        try {
            withContext(Dispatchers.IO) {
                TerminalLogger.debugf { "connecting to QUIC $quicHost:$quicPort" }
                val builder = QuicClientConnection.newBuilder()
                    .uri(URI("https://$quicHost:$quicPort"))
                    .applicationProtocol("punchnet/1.0")
                    .noServerCertificateCheck()
                connection = builder.build()
                connection!!.connect()
                
                // Open bidirectional stream
                stream = connection!!.createStream(true)
                
                // In QUIC, streams are created implicitly when the first data is sent.
                // We send a PACKET_EMPTY (0x00) so the server can accept the stream and send Welcome.
                val emptyPacket = ByteArray(3)
                emptyPacket[0] = 0 // Len high
                emptyPacket[1] = 1 // Len low (1 byte for PacketType)
                emptyPacket[2] = 0 // PacketType.Empty (0x00)
                TerminalLogger.debugf { "Sending QUIC application packet (PACKET_EMPTY): [00 01 00]" }
                stream!!.outputStream.write(emptyPacket)
                stream!!.outputStream.flush()
            }
        } catch (e: Exception) {
            TerminalLogger.errorf { "failed to connect to QUIC ${quicHost}:${quicPort}: $e" }
            delay(3000)
            continue
        }

        try {
            val job_read_packet = scope.launch(Dispatchers.IO) {
                try {
                    while (isActive) {
                        val controlPacket = readQUICPacket(stream!!)
                        if (controlPacket == null) {
                            TerminalLogger.errorf { "QUIC Stream EOF or error" }
                            break
                        }
                        onMessage(scope, controlPacket)
                    }
                } finally {
                    TerminalLogger.errorf { "QUIC read loop closing" }
                }
            }

            val job_write_to_packet = scope.launch(Dispatchers.IO) {
                try {
                    while (isActive) {
                        val msg = toSocket.receive()
                        val packetType = if (msg.size >= 3) PacketType.fromValue(msg[2].toUByte()) else null
                        TerminalLogger.debugf { "Sending QUIC application packet: type=${packetType ?: "unknown"}, bytes=${msg.size}" }
                        stream!!.outputStream.write(msg)
                    }
                } catch (e: Exception) {
                    TerminalLogger.errorf { "failed to write to QUIC stream: $e" }
                }
            }

            val job_check_pong = scope.launch(Dispatchers.IO) {
                while (isActive) {
                    delay(10_000)
                    val now = System.currentTimeMillis() / 1000
                    if (connected.get() && (now - pong_time.get() > com.jihe.punchnet.sdlan.config.SDLanConfig.TCP_PING_TIME * 2)) {
                        TerminalLogger.errorf { "QUIC pong check expired" }
                        break
                    }
                }
            }

            val job_check_stop = scope.launch(Dispatchers.IO) {
                while (isActive) {
                    try {
                        val v = start_stop.receive()
                        if (v.flag == StartStopFlag.IsStop) {
                            started.set(false)
                            break
                        }
                    } catch (e: Exception) {
                        break
                    }
                }
            }

            TerminalLogger.debugf { "QUIC Stream connected" }
            onConnected(stream!!)
            connectingChan?.send(ConnectingState.Connected)

            select<Unit> {
                job_read_packet.onJoin { }
                job_write_to_packet.onJoin { }
                job_check_pong.onJoin { }
                job_check_stop.onJoin { }
            }

            job_read_packet.cancelAndJoin()
            job_write_to_packet.cancelAndJoin()
            job_check_pong.cancelAndJoin()
            job_check_stop.cancelAndJoin()

            delay(1000)
        } catch (e: Exception) {
            onDisconnected()
        } finally {
            try {
                connection?.close()
            } catch (e: Exception) {}
        }
    }
}

suspend fun readQUICPacket(stream: QuicStream): SDLanControl? = withContext(Dispatchers.IO) {
    try {
        // <<Len:16, PacketType:8, ProtobufPayload>>
        val lenBytes = ByteArray(2)
        if (!stream.inputStream.readFullyOrEof(lenBytes)) return@withContext null
        
        val len = ByteBuffer.wrap(lenBytes).short.toInt() and 0xFFFF
        if (len < 1) return@withContext null
        
        val typeByte = ByteArray(1)
        if (!stream.inputStream.readFullyOrEof(typeByte)) return@withContext null
        
        val packetTypeUByte = typeByte[0].toUByte()
        val packetType = PacketType.fromValue(packetTypeUByte)
        if (packetType == null) {
            TerminalLogger.errorf { "packet type not found: $packetTypeUByte" }
            return@withContext null
        }
        
        val payloadLen = len - 1
        val payload = ByteArray(payloadLen)
        var toread = payloadLen
        var offset = 0
        while (toread > 0) {
            val sizeGot = stream.inputStream.read(payload, offset, toread)
            if (sizeGot <= 0) return@withContext null
            toread -= sizeGot
            offset += sizeGot
        }
        
        return@withContext SDLanControl(packetType, payload)
    } catch (e: Exception) {
        TerminalLogger.errorf { "failed to read QUIC stream: $e" }
        return@withContext null
    }
}

private fun InputStream.readFullyOrEof(buffer: ByteArray): Boolean {
    var offset = 0
    while (offset < buffer.size) {
        val readSize = read(buffer, offset, buffer.size - offset)
        if (readSize <= 0) return false
        offset += readSize
    }
    return true
}
