import asyncio
import logging
import signal
import sys
from aioquic.asyncio import QuicConnectionProtocol, serve
from aioquic.quic.configuration import QuicConfiguration
from aioquic.quic.events import ConnectionTerminated, StreamDataReceived, StreamReset
import parameters

# Setup basic logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("QuicServer")

# Global list to track active protocol instances
active_protocols = []

class TunnelServerProtocol(QuicConnectionProtocol):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.loop = asyncio.get_event_loop()
        self.tcp_connections = {}  # Map TCP connections to QUIC streams
        self.udp_connections = {}  # Map UDP connections to QUIC streams
        self.udp_last_activity = {}  # Track last activity time for UDP connections
        active_protocols.append(self)  # Add this protocol instance to the list
        try:
            asyncio.create_task(self.cleanup_stale_udp_connections())
        except Exception as e:
            logger.info(f"Error in cleanup_stale_udp task: {e}")

    def connection_lost(self, exc):
        logger.info("Quic channel lost")
        if self in active_protocols:
            active_protocols.remove(self)
        super().connection_lost(exc)
        self.close_all_tcp_connections()
        self.close_all_udp_connections()

    def close_all_tcp_connections(self):
        logger.info("Closing all TCP connections from server...")
        for stream_id, (reader, writer) in self.tcp_connections.items():
            logger.info(f"Closing TCP connection for stream {stream_id}...")
            writer.close()
        self.tcp_connections.clear()

    def close_all_udp_connections(self):
        logger.info("Closing all UDP connections from server...")
        for stream_id, (transport, _) in self.udp_connections.items():
            logger.info(f"Closing UDP connection for stream {stream_id}...")
            transport.close()
        self.udp_connections.clear()
        self.udp_last_activity.clear()


    def close_this_stream(self, stream_id):
        try:
            logger.info(f"FIN to stream={stream_id} sent")
            self._quic.send_stream_data(stream_id, b"", end_stream=True)  # Send FIN flag
            self.transmit()  # Send the FIN flag over the network
        except Exception as e:
            logger.info(f"Error closing stream at server: {e}")

        try:
            if stream_id in self.tcp_connections:
                writer = self.tcp_connections[stream_id][1]
                writer.close()
                del self.tcp_connections[stream_id]
            if stream_id in self.udp_connections:
                transport, _ = self.udp_connections[stream_id]
                transport.close()
                del self.udp_connections[stream_id]
                del self.udp_last_activity[stream_id]
        except Exception as e:
            logger.info(f"Error closing socket at server: {e}")




    async def cleanup_stale_udp_connections(self):
        logger.info("UDP cleanup task running!")
        check_time = min(parameters.udp_timeout,60)
        while True:
            await asyncio.sleep(check_time)  # Run cleanup every 60 seconds
            current_time = self.loop.time()
            stale_streams = [
                stream_id for stream_id, last_time in self.udp_last_activity.items()
                if current_time - last_time > parameters.udp_timeout
            ]
            for stream_id in stale_streams:
                logger.info(f"idle UDP stream={stream_id} timeout reached")
                self.close_this_stream(stream_id)



    async def forward_tcp_to_quic(self, stream_id, reader):
        logger.info(f"Task TCP to QUIC started")
        try:
            while True:
                data = await reader.read(4096)  # Read data from TCP socket
                if not data:
                    break
                # logger.info(f"Forwarding data from TCP to QUIC on stream {stream_id}")
                self._quic.send_stream_data(stream_id=stream_id, data=data, end_stream=False)
                self.transmit()  # Flush
        except Exception as e:
            logger.info(f"Error forwarding TCP to QUIC: {e}")
        finally:
            logger.info(f"Task TCP to QUIC Ended")
            self.close_this_stream(stream_id)



    async def connect_tcp(self, stream_id, target_port):
        logger.info(f"Connecting to TCP:{target_port}...")
        try:
            reader, writer = await asyncio.open_connection(parameters.xray_server_ip_address, target_port)
            logger.info(f"TCP connection established for stream {stream_id} to port {target_port}")

            # Start forwarding data from TCP to QUIC
            asyncio.create_task(self.forward_tcp_to_quic(stream_id, reader))

            resp_data = parameters.quic_auth_code + "i am ready,!###!"
            self._quic.send_stream_data(stream_id=stream_id, data=resp_data.encode("utf-8"), end_stream=False)
            self.transmit()  # Flush

            self.tcp_connections[stream_id] = (reader, writer)
        except Exception as e:
            logger.info(f"Failed to establish TCP:{target_port} connection: {e}")
            self.close_this_stream(stream_id)



    async def forward_udp_to_quic(self, stream_id, protocol):
        logger.info(f"Task UDP to QUIC started")
        try:
            while True:
                data, _ = await protocol.queue.get()  # Wait for data from UDP
                if(data == None):
                    break
                # logger.info(f"Forwarding data from UDP to QUIC on stream {stream_id}")
                self._quic.send_stream_data(stream_id, data)
                self.transmit()  # Flush
                self.udp_last_activity[stream_id] = self.loop.time()
        except Exception as e:
            logger.info(f"Error forwarding UDP to QUIC: {e}")
        finally:
            logger.info(f"Task UDP to QUIC Ended")
            self.close_this_stream(stream_id)


    async def connect_udp(self, stream_id, target_port):
        class UdpProtocol:
            def __init__(self):
                self.transport = None
                self.queue = asyncio.Queue()
                self.stream_id = stream_id

            def connection_made(self, transport):
                self.transport = transport

            def datagram_received(self, data, addr):
                logger.info(f"put this to queue data={data} addr={addr}")
                self.queue.put_nowait((data, addr))

            def error_received(self, exc):
                logger.info(f"UDP error received: {exc}")
                self.queue.put_nowait((None, None)) # to cancel task
                if(self.transport):
                        self.transport.close()
                        logger.info("UDP transport closed")

            def connection_lost(self, exc):
                logger.info("UDP connection lost.")
                self.queue.put_nowait((None, None)) # to cancel task
                if(self.transport):
                        self.transport.close()
                        logger.info("UDP transport closed")

        try:
            # Create a UDP socket
            logger.info(f"Connecting to UDP:{target_port}...")
            loop = asyncio.get_event_loop()
            transport, protocol = await loop.create_datagram_endpoint(
                UdpProtocol,
                remote_addr=(parameters.xray_server_ip_address, target_port)
            )
            self.udp_connections[stream_id] = (transport, protocol)
            self.udp_last_activity[stream_id] = self.loop.time()  # Track last activity time
            logger.info(f"UDP connection established for stream {stream_id} to port {target_port}")

            asyncio.create_task(self.forward_udp_to_quic(stream_id, protocol))
        except Exception as e:
            logger.info(f"Failed to establish UDP connection: {e}")




    def quic_event_received(self, event):
        # print("EVENT",event)
        if isinstance(event, StreamDataReceived):
            try:
                # logger.info(f"Server received from QUIC on stream {event.stream_id}")
                # logger.info(f"Server TCP IDs -> {self.tcp_connections.keys()}")
                # logger.info(f"Server UDP IDs -> {self.udp_connections.keys()}")

                if event.end_stream:
                    logger.info(f"Stream={event.stream_id} closed by client.")
                    self.close_this_stream(event.stream_id)

                # Forward data to the corresponding TCP connection
                elif event.stream_id in self.tcp_connections:
                    writer = self.tcp_connections[event.stream_id][1]
                    writer.write(event.data)  # Send data over TCP
                    try:
                        asyncio.create_task(writer.drain())
                    except ConnectionResetError as e42:
                        logger.info(f"ERR in writer drain task : {e42}")
                    except Exception as e43:
                        logger.info(f"ERR in writer drain task : {e43}")

                # Forward data to the corresponding UDP connection
                elif event.stream_id in self.udp_connections:
                    transport, _ = self.udp_connections[event.stream_id]
                    transport.sendto(event.data)  # Send data over UDP
                    self.udp_last_activity[event.stream_id] = self.loop.time()  # Update last activity time

                else:
                    socket_type = None
                    socket_port = 0

                    # Assume req is like => auth+"connect,udp,443,!###!"
                    new_req = event.data.split(b",!###!", 1)
                    req_header = ""
                    try:
                        req_header = new_req[0].decode("utf-8")
                    except Exception as e47:
                        logger.info(f"ERR in req decoding : {e47}")
                        req_header=""

                    logger.info("New req comes -> " + req_header)

                    if req_header.startswith(parameters.quic_auth_code + "connect,"):
                        j = len(parameters.quic_auth_code) + 8
                        if req_header[j:j + 3] == "tcp":
                            socket_type = "tcp"
                        elif req_header[j:j + 3] == "udp":
                            socket_type = "udp"

                        try:
                            socket_port = int(req_header[j + 4:])
                        except ValueError:
                            logger.info("Invalid port.")

                        if socket_port > 0:
                            if socket_type == "tcp":
                                # New stream detected, create a TCP connection
                                asyncio.create_task(self.connect_tcp(event.stream_id, socket_port))
                            elif socket_type == "udp":
                                # New stream detected, create a UDP connection
                                asyncio.create_task(self.connect_udp(event.stream_id, socket_port))
                            else:
                                logger.info("Invalid Req: socket type unknown.")
                        else:
                            logger.info("Invalid Req: socket port unknown.")
                    else:
                        logger.info("Invalid Req header")

            except Exception as e:
                logger.info(f"Quic event server error: {e}")

        elif isinstance(event, StreamReset):
            # Handle stream reset (client closed the stream)
            logger.info(f"Stream {event.stream_id} reset by client.")
            self.close_this_stream(event.stream_id)

        elif isinstance(event, ConnectionTerminated):
            logger.info(f"Connection lost: {event.reason_phrase}")
            self.connection_lost(event.reason_phrase)


async def run_server():
    configuration = QuicConfiguration(is_client=False)
    configuration.load_cert_chain(parameters.quic_cert_filepath[0], parameters.quic_cert_filepath[1])
    configuration.max_data = parameters.quic_max_data
    configuration.max_stream_data = parameters.quic_max_stream_data
    configuration.idle_timeout = parameters.quic_idle_timeout
    configuration.max_datagram_size = parameters.quic_mtu

    # Start QUIC server
    await serve("0.0.0.0", parameters.quic_server_port, configuration=configuration, create_protocol=TunnelServerProtocol)
    logger.warning(f"Server listening for QUIC on port {parameters.quic_server_port}")

    # Keep the server running
    await asyncio.Future()  # Run forever


def handle_shutdown(signum, frame):
    logger.info("Shutting down server gracefully...")
    for protocol in active_protocols:
        protocol.close_all_tcp_connections()
        protocol.close_all_udp_connections()
        protocol.close()
    logger.info("Server shutdown complete.")
    sys.exit(0)


if __name__ == "__main__":
    signal.signal(signal.SIGTERM, handle_shutdown)
    signal.signal(signal.SIGINT, handle_shutdown)

    asyncio.run(run_server())
