import asyncio
import sys


class LoadBalancer:
    def __init__(self, backend_ports):
        self.backend_ports = backend_ports
        self.backend_processes = []
        self.client_counts = [0] * len(backend_ports)
        self.lock = asyncio.Lock()

    async def start_backend_servers(self, port, workers):
        for x in range(workers):
            port += 1
            process = await asyncio.create_subprocess_exec(
                sys.executable,
                sys.argv[0],
                "backend",
                str(port),
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            port += 1
            self.backend_processes.append(process)
            print(
                f"Started backend server on port {(port-1)/port} with PID {process.pid}"
            )

    async def handle_client(self, reader, writer):
        async with self.lock:
            min_clients = min(self.client_counts)
            server_index = self.client_counts.index(min_clients)
            self.client_counts[server_index] += 1
        backend = ("127.0.0.1", self.backend_ports[server_index])
        try:
            backend_reader, backend_writer = await asyncio.open_connection(*backend)

            async def forward(r, w):
                try:
                    while True:
                        data = await r.read(1024)
                        if not data:
                            break
                        w.write(data)
                        await w.drain()
                except asyncio.CancelledError:
                    pass
                finally:
                    w.close()

            task1 = asyncio.create_task(forward(reader, backend_writer))
            task2 = asyncio.create_task(forward(backend_reader, writer))
            await asyncio.gather(task1, task2)
        except Exception as e:
            print(f"Error: {e}")
        finally:
            writer.close()
            async with self.lock:
                self.client_counts[server_index] -= 1

    async def monitor(self):
        while True:
            await asyncio.sleep(5)
            print("Connected clients per server:")
            for i, count in enumerate(self.client_counts):
                print(f"Server {self.backend_ports[i]}: {count} clients")

    async def start(self, host="0.0.0.0", port=8081, workers=5):
        await self.start_backend_servers(port, workers)
        server = await asyncio.start_server(self.handle_client, host, port)
        asyncio.create_task(self.monitor())

        # Handle shutdown gracefully
        try:
            async with server:
                await server.serve_forever()
        except asyncio.CancelledError:
            pass
        finally:
            # Terminate backend processes
            for process in self.backend_processes:
                process.terminate()
            await asyncio.gather(*(p.wait() for p in self.backend_processes))
            print("Backend processes terminated.")


async def backend_echo_server(port):
    async def handle_echo(reader, writer):
        try:
            while True:
                data = await reader.read(1024)
                if not data:
                    break
                writer.write(data)
                await writer.drain()
        except Exception:
            pass
        finally:
            writer.close()

    server = await asyncio.start_server(handle_echo, "127.0.0.1", port)
    print(f"Backend echo server running on port {port}")
    await server.serve_forever()


async def main():
    backend_ports = [8001, 8003, 8005, 8006]
    # Launch backend echo servers
    # Wait a moment for servers to start
    lb = LoadBalancer(backend_ports)
    await lb.start()


if __name__ == "__main__":
    if len(sys.argv) > 1:
        if sys.argv[1] == "backend":
            port = int(sys.argv[2])
            from snek.app import Application

            snek = Application(port=port)
            web.run_app(snek, port=port, host="127.0.0.1")
        elif sys.argv[1] == "sync":

            web.run_app(snek, port=port, host="127.0.0.1")
    else:
        try:
            asyncio.run(main())
        except KeyboardInterrupt:
            print("Shutting down...")
