import asyncio
import os
import websockets
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# The remote HF Space websocket URL (e.g., wss://augment17-redis.hf.space/ws)
REMOTE_WS_URL = os.getenv("REMOTE_REDIS_WS_URL")
LOCAL_PORT = int(os.getenv("LOCAL_REDIS_PORT", 6379))

async def handle_client(reader, writer):
    if not REMOTE_WS_URL:
        logger.error("REMOTE_REDIS_WS_URL is not set!")
        writer.close()
        return

    logger.info("New local TCP connection. Connecting to remote WebSocket...")
    try:
        async with websockets.connect(REMOTE_WS_URL) as ws:
            logger.info("Connected to remote WebSocket.")

            async def pipe_tcp_to_ws():
                try:
                    while True:
                        data = await reader.read(4096)
                        if not data:
                            break
                        await ws.send(data)
                except Exception as e:
                    logger.info(f"tcp_to_ws closed: {e}")

            async def pipe_ws_to_tcp():
                try:
                    while True:
                        data = await ws.recv()
                        writer.write(data)
                        await writer.drain()
                except Exception as e:
                    logger.info(f"ws_to_tcp closed: {e}")

            await asyncio.gather(
                pipe_tcp_to_ws(),
                pipe_ws_to_tcp(),
                return_exceptions=True
            )
    except Exception as e:
        logger.error(f"WebSocket connection failed: {e}")
    finally:
        writer.close()
        logger.info("Local TCP connection closed.")

async def main():
    server = await asyncio.start_server(handle_client, '127.0.0.1', LOCAL_PORT)
    addr = server.sockets[0].getsockname()
    logger.info(f"Local WebSocket-to-TCP proxy listening on {addr}")

    async with server:
        await server.serve_forever()

if __name__ == '__main__':
    asyncio.run(main())