"""
Bedrock Edition Tunnel — Client Side

Listens on a local UDP port (default 19132) and relays all packets
through a WebSocket to the server-side UDP<->WS bridge, which
forwards them to Geyser.
"""

import asyncio
import threading
import logging

try:
    import websockets
except ImportError:
    raise SystemExit("Install websockets: pip install websockets")

logger = logging.getLogger("bedrock-tunnel")


class _LocalUDP(asyncio.DatagramProtocol):
    """Receives datagrams from the local Bedrock client."""

    def __init__(self):
        self.transport = None
        self.ws = None
        self.client_addr = None

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

    def datagram_received(self, data, addr):
        self.client_addr = addr
        if self.ws is not None:
            asyncio.ensure_future(self._forward(data))

    async def _forward(self, data):
        try:
            await self.ws.send(data)
        except Exception:
            pass

    def send_to_client(self, data: bytes):
        if self.client_addr and self.transport:
            self.transport.sendto(data, self.client_addr)


class BedrockTunnel:
    """Manages the local UDP listener and the WS relay to the server bridge."""

    def __init__(self, on_log=None, on_status=None, on_port=None):
        self.on_log = on_log or (lambda _: None)
        self.on_status = on_status or (lambda _: None)
        self.on_port = on_port or (lambda _: None)
        self.running = False
        self._thread: threading.Thread | None = None
        self._loop: asyncio.AbstractEventLoop | None = None
        self._ws = None
        self._protocol = None
        self.local_port = 0

    def start(self, hostname: str, local_port: int):
        """Start the Bedrock tunnel."""
        if self.running:
            self.on_log("Bedrock tunnel already running")
            return
        self.running = True
        self.local_port = local_port
        self._thread = threading.Thread(
            target=self._run, args=(hostname, local_port), daemon=True
        )
        self._thread.start()

    def stop(self):
        """Stop the Bedrock tunnel."""
        if not self.running:
            return
        self.running = False
        if self._loop and not self._loop.is_closed():
            self._loop.call_soon_threadsafe(self._shutdown)
        if self._thread:
            self._thread.join(timeout=5)

    def _run(self, hostname: str, local_port: int):
        """Thread entry point."""
        self._loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._loop)
        try:
            self._loop.run_until_complete(self._main(hostname, local_port))
        except Exception as exc:
            self.on_log(f"Bedrock tunnel error: {exc}")
        finally:
            self._loop.close()

    async def _main(self, hostname: str, local_port: int):
        """Main async coroutine."""
        self.on_status("connecting")
        
        # Start local UDP server
        loop = asyncio.get_event_loop()
        try:
            transport, protocol = await loop.create_datagram_endpoint(
                lambda: _LocalUDP(), local_addr=("127.0.0.1", local_port)
            )
            self._protocol = protocol
            actual_port = transport.get_extra_info('sockname')[1]
            self.local_port = actual_port
            self.on_port(actual_port)
            self.on_log(f"Bedrock: Listening on localhost:{actual_port}")
        except OSError as e:
            self.on_log(f"Bedrock: Cannot bind port {local_port}: {e}")
            self.on_status("error")
            self.running = False
            return

        # Connect to WebSocket
        ws_url = f"wss://{hostname}"
        retry_count = 0
        
        while self.running:
            try:
                async with websockets.connect(
                    ws_url,
                    max_size=None,
                    ping_interval=20,
                    ping_timeout=30,
                ) as ws:
                    self._ws = ws
                    self._protocol.ws = ws
                    self.on_status("connected")
                    self.on_log(f"Bedrock: Connected to {hostname}")
                    retry_count = 0
                    
                    # Relay WebSocket → UDP
                    async for message in ws:
                        if not self.running:
                            break
                        if isinstance(message, bytes):
                            self._protocol.send_to_client(message)
                        
            except (websockets.ConnectionClosed, OSError) as e:
                self._ws = None
                self._protocol.ws = None
                if self.running:
                    retry_count += 1
                    wait = min(retry_count * 2, 30)
                    self.on_status("reconnecting")
                    self.on_log(f"Bedrock: Connection lost, retry in {wait}s")
                    await asyncio.sleep(wait)
                    
        self.on_status("disconnected")
        transport.close()

    def _shutdown(self):
        """Shutdown callback for loop."""
        if self._ws:
            asyncio.ensure_future(self._ws.close())
