#!/usr/bin/env python3
"""
Local TCP -> WebSocket proxy for connecting to Minecraft on HuggingFace.

Usage:
  python3 client-proxy.py <your-space-url>

Example:
  python3 client-proxy.py https://username-minecraft-server.hf.space

Then connect your Minecraft client to: localhost:25565
"""

import asyncio
import sys
import websockets
import socket

HF_WS_PORT = 25565  # Local port to listen on
BUFFER_SIZE = 4096


async def bridge_tcp_to_ws(reader, ws):
    """Forward data from local TCP client to remote WebSocket."""
    try:
        while True:
            data = await reader.read(BUFFER_SIZE)
            if not data:
                break
            await ws.send(data)
    except Exception as e:
        print(f"[TCP->WS] Error: {e}")
    finally:
        print("[TCP->WS] Connection closed")


async def bridge_ws_to_tcp(ws, writer):
    """Forward data from remote WebSocket to local TCP client."""
    try:
        async for message in ws:
            writer.write(message)
            await writer.drain()
    except Exception as e:
        print(f"[WS->TCP] Error: {e}")
    finally:
        print("[WS->TCP] Connection closed")


async def handle_client(reader, writer, ws_url):
    """Handle a single Minecraft client connection."""
    print(f"[NEW] Client connected, bridging to {ws_url}")
    try:
        async with websockets.connect(ws_url) as ws:
            await asyncio.gather(
                bridge_tcp_to_ws(reader, ws),
                bridge_ws_to_tcp(ws, writer),
            )
    except Exception as e:
        print(f"[ERROR] WebSocket connection failed: {e}")
    finally:
        writer.close()
        await writer.wait_closed()
        print("[CLOSED] Client disconnected")


async def main(ws_url):
    """Start local TCP server that proxies to WebSocket."""
    server = await asyncio.start_server(
        lambda r, w: handle_client(r, w, ws_url),
        "127.0.0.1",
        HF_WS_PORT,
    )
    print(f"Local proxy listening on 127.0.0.1:{HF_WS_PORT}")
    print(f"Forwarding to: {ws_url}")
    print(f"Connect your Minecraft client to: localhost:{HF_WS_PORT}")
    async with server:
        await server.serve_forever()


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <hf-space-url>")
        print(f"Example: {sys.argv[0]} https://user-minecraft.hf.space")
        sys.exit(1)

    url = sys.argv[1].rstrip("/")
    ws_url = url.replace("https://", "wss://").replace("http://", "ws://")

    # Install websockets if needed
    try:
        import websockets
    except ImportError:
        import subprocess
        print("Installing websockets library...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "websockets"])
        import websockets

    asyncio.run(main(ws_url))