"""
Java Edition Tunnel — Client Side

Keeps a persistent WebSocket to the server bridge. Local TCP connections
from Minecraft are relayed over the single WS with minimal framing:
  0x01 + data = new TCP session (+ first chunk)
  0x02 + data = data for current session
  0x03        = close current session
"""

import asyncio
import struct
import threading
import logging

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

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

MSG_NEW   = b'\x01'
MSG_DATA  = b'\x02'
MSG_CLOSE = b'\x03'


class JavaTunnel:
    """Persistent-WS Java tunnel with one-at-a-time TCP sessions."""

    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._server: asyncio.Server | None = None
        self._ws = None
        self._ws_lock = None
        self._active_writer = None
        self.local_port = 0

    def start(self, hostname: str, local_port: int):
        """Start the Java tunnel."""
        if self.running:
            self.on_log("Java 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 Java 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"Java tunnel error: {exc}")
        finally:
            self._loop.close()

    async def _main(self, hostname: str, local_port: int):
        """Main async coroutine."""
        self._ws_lock = asyncio.Lock()
        self.on_status("connecting")
        
        # Start local TCP server
        try:
            self._server = await asyncio.start_server(
                self._handle_local_tcp, "127.0.0.1", local_port
            )
            actual_port = self._server.sockets[0].getsockname()[1]
            self.local_port = actual_port
            self.on_port(actual_port)
            self.on_log(f"Java: Listening on localhost:{actual_port}")
        except OSError as e:
            self.on_log(f"Java: 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.on_status("connected")
                    self.on_log(f"Java: Connected to {hostname}")
                    retry_count = 0
                    
                    # Read from WebSocket
                    async for message in ws:
                        if not self.running:
                            break
                        await self._handle_ws_message(message)
                        
            except (websockets.ConnectionClosed, OSError) as e:
                self._ws = None
                if self.running:
                    retry_count += 1
                    wait = min(retry_count * 2, 30)
                    self.on_status("reconnecting")
                    self.on_log(f"Java: Connection lost, retry in {wait}s")
                    await asyncio.sleep(wait)
                    
        self.on_status("disconnected")

    async def _handle_local_tcp(self, reader, writer):
        """Handle incoming local TCP connection from Minecraft client."""
        try:
            # Close any existing session
            if self._active_writer and not self._active_writer.is_closing():
                self._active_writer.close()
                await self._active_writer.wait_closed()
                async with self._ws_lock:
                    if self._ws:
                        await self._ws.send(MSG_CLOSE)
            
            self._active_writer = writer
            
            # Read first chunk and send as NEW
            data = await reader.read(65536)
            if not data:
                writer.close()
                await writer.wait_closed()
                return
                
            async with self._ws_lock:
                if self._ws:
                    await self._ws.send(MSG_NEW + data)
            
            # Relay subsequent chunks
            while self.running:
                data = await reader.read(65536)
                if not data:
                    break
                async with self._ws_lock:
                    if self._ws:
                        await self._ws.send(MSG_DATA + data)
                        
        except Exception as exc:
            self.on_log(f"Java TCP error: {exc}")
        finally:
            if not writer.is_closing():
                writer.close()
                await writer.wait_closed()
            async with self._ws_lock:
                if self._ws and self._active_writer == writer:
                    await self._ws.send(MSG_CLOSE)
                    self._active_writer = None

    async def _handle_ws_message(self, message):
        """Handle message from WebSocket."""
        if isinstance(message, str):
            message = message.encode("utf-8")
        if not message:
            return
            
        msg_type = message[0:1]
        payload = message[1:]
        
        if msg_type == MSG_DATA and self._active_writer and not self._active_writer.is_closing():
            try:
                self._active_writer.write(payload)
                await self._active_writer.drain()
            except Exception:
                pass
        elif msg_type == MSG_CLOSE and self._active_writer:
            if not self._active_writer.is_closing():
                self._active_writer.close()
                await self._active_writer.wait_closed()
            self._active_writer = None

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