"""mldash connect -- WebSocket bridge for ML-Dash Studio.

Runs a local WebSocket server that receives code from ML-Dash Studio
and forwards it to a persistent Python kernel, streaming output back.

Architecture:
  Browser <-> WebSocket <-> Bridge <-> Kernel (stdin/stdout)

The kernel is a long-running Python process (_kernel.py) that maintains
a persistent namespace. Variables survive between cell executions.
"""

import asyncio
import json
import os
import signal
import secrets
import sys
import time
from urllib.parse import urlparse, parse_qs

try:
    import websockets
except ImportError:
    print("Error: websockets is required. Install: pip install 'websockets>=12.0'")
    sys.exit(1)

import requests

from ._credentials import load as _load_credentials
from ._env import resolve_kernel_python

# Magic delimiter used by the kernel to prefix JSON protocol messages
KERNEL_DELIMITER = "\x00MLDASH_MSG\x00"


def _init_autocomplete():
    """Try to initialize the autocomplete engine. Returns engine or None."""
    try:
        from ._autocomplete import AutocompleteEngine, model_exists
        if not model_exists():
            return None
        engine = AutocompleteEngine()
        engine.load_in_background()
        return engine
    except Exception:
        return None

def start_bridge(port=0, url=None, api_key=None):
    """Start the local compute bridge."""
    # Show which Python the kernel will use
    kernel_python, python_source = resolve_kernel_python()
    source_labels = {
        "env": f"Using $MLDASH_PYTHON: {kernel_python}",
        "managed": f"Using managed environment: {kernel_python}",
        "system": f"Using system Python: {kernel_python}",
    }
    print(source_labels[python_source])

    # Resolve credentials
    creds = _load_credentials()
    url = url or creds.get('url') or os.environ.get('MLDASH_URL')
    api_key = api_key or creds.get('api_key') or os.environ.get('MLDASH_API_KEY')

    if not url or not api_key:
        print("Not logged in. Run `mldash login` first.")
        sys.exit(1)

    # Generate session token
    token = secrets.token_urlsafe(32)

    # Auto-select port if not specified
    if port == 0:
        import socket
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(('localhost', 0))
            port = s.getsockname()[1]

    # Start the async event loop
    try:
        asyncio.run(_run_bridge(port, token, url, api_key))
    except KeyboardInterrupt:
        print("\nBridge stopped.")


class _KernelManager:
    """Manages a persistent Python kernel subprocess."""

    def __init__(self):
        self.process = None
        self._reader_task = None
        self._stderr_task = None
        self._websocket = None
        self._current_cell_id = None

    async def start(self, websocket=None):
        """Start (or restart) the kernel process."""
        await self.stop()

        kernel_python, _ = resolve_kernel_python()

        self.process = await asyncio.create_subprocess_exec(
            kernel_python, "-m", "mldash._kernel",
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        if websocket:
            self._websocket = websocket

        self._reader_task = asyncio.create_task(self._read_stdout())
        self._stderr_task = asyncio.create_task(self._read_stderr())
        print("Kernel started (pid={}).".format(self.process.pid))

    async def stop(self):
        """Stop the kernel process gracefully."""
        if self._reader_task:
            self._reader_task.cancel()
            self._reader_task = None
        if self._stderr_task:
            self._stderr_task.cancel()
            self._stderr_task = None

        if self.process and self.process.returncode is None:
            self.process.terminate()
            try:
                await asyncio.wait_for(self.process.wait(), timeout=5)
            except asyncio.TimeoutError:
                self.process.kill()
            print("Kernel stopped.")

        self.process = None

    @property
    def alive(self):
        return self.process is not None and self.process.returncode is None

    def set_websocket(self, ws):
        self._websocket = ws

    async def execute(self, code, cell_id, execution_id):
        """Send code to the kernel for execution."""
        if not self.alive:
            await self.start(self._websocket)

        self._current_cell_id = cell_id
        cmd = json.dumps({
            "type": "execute",
            "code": code,
            "cell_id": cell_id,
            "execution_id": execution_id,
        }) + "\n"

        self.process.stdin.write(cmd.encode())
        await self.process.stdin.drain()

    async def restart(self):
        """Restart kernel state (clears namespace)."""
        if self.alive:
            # Send restart command — kernel clears namespace and sends kernel_ready
            cmd = json.dumps({"type": "restart"}) + "\n"
            self.process.stdin.write(cmd.encode())
            await self.process.stdin.drain()
        else:
            await self.start(self._websocket)

    async def interrupt(self):
        """Send SIGINT to interrupt the current cell execution."""
        if self.alive:
            try:
                self.process.send_signal(signal.SIGINT)
            except ProcessLookupError:
                pass

    async def _read_stdout(self):
        """Read kernel stdout, parse protocol messages, forward to browser."""
        try:
            while True:
                raw = await self.process.stdout.readline()
                if not raw:
                    break  # kernel exited

                text = raw.decode("utf-8", errors="replace").rstrip("\n")

                if text.startswith(KERNEL_DELIMITER):
                    # Protocol message — parse JSON and forward to browser
                    json_str = text[len(KERNEL_DELIMITER):]
                    try:
                        msg = json.loads(json_str)
                        if self._websocket:
                            await self._websocket.send(json.dumps(msg))
                    except (json.JSONDecodeError, Exception):
                        pass
                elif text and self._websocket and self._current_cell_id:
                    # Raw output (from C extensions, os.write, etc.)
                    try:
                        await self._websocket.send(json.dumps({
                            "type": "stdout",
                            "data": text + "\n",
                            "cell_id": self._current_cell_id,
                            "execution_id": "",
                        }))
                    except Exception:
                        pass

        except asyncio.CancelledError:
            pass
        except Exception as e:
            print(f"Kernel stdout reader error: {e}")

    async def _read_stderr(self):
        """Read kernel stderr for kernel-level errors (not user code errors)."""
        try:
            while True:
                raw = await self.process.stderr.readline()
                if not raw:
                    break
                text = raw.decode("utf-8", errors="replace").rstrip()
                if text:
                    print(f"[kernel] {text}")
        except asyncio.CancelledError:
            pass
        except Exception:
            pass


async def _handle_complete(ws, engine, msg):
    """Handle a completion request from Studio."""
    prefix = msg.get("prefix", "")
    suffix = msg.get("suffix", "")
    request_id = msg.get("request_id", 0)
    cell_id = msg.get("cell_id", "")

    text = await engine.complete(prefix, suffix, request_id)
    if text:
        await ws.send(json.dumps({
            "type": "completion",
            "text": text,
            "request_id": request_id,
            "cell_id": cell_id,
        }))


async def _run_bridge(port, token, url, api_key):
    """Run the WebSocket bridge server."""
    kernel = _KernelManager()
    autocomplete_engine = _init_autocomplete()

    async def handler(websocket, path=None):
        # Extract token from query params (compat: websockets 10-13+)
        request_path = path
        if request_path is None:
            try:
                request_path = websocket.request.path
            except AttributeError:
                try:
                    request_path = websocket.path
                except AttributeError:
                    request_path = '/'

        parsed = parse_qs(urlparse(request_path).query)
        client_token = parsed.get('token', [None])[0]

        if client_token != token:
            print(f"Connection rejected: invalid token")
            await websocket.close(4001, "Invalid token")
            return

        print("Studio connected.")
        kernel.set_websocket(websocket)

        # Start kernel if not already running
        if not kernel.alive:
            await kernel.start(websocket)

        try:
            async for raw_message in websocket:
                try:
                    msg = json.loads(raw_message)
                except json.JSONDecodeError:
                    continue

                msg_type = msg.get("type")

                if msg_type == "ping":
                    await websocket.send(json.dumps({"type": "pong"}))

                elif msg_type == "execute":
                    code = msg.get("code", "")
                    cell_id = msg.get("cell_id", "")
                    execution_id = msg.get("execution_id", "")
                    print(f"Executing cell {cell_id[:8]}... ({len(code)} chars)")
                    await kernel.execute(code, cell_id, execution_id)

                elif msg_type == "interrupt":
                    print("Interrupting execution...")
                    await kernel.interrupt()

                elif msg_type == "kernel_restart":
                    print("Restarting kernel...")
                    await kernel.restart()

                elif msg_type == "complete":
                    if autocomplete_engine:
                        asyncio.create_task(
                            _handle_complete(websocket, autocomplete_engine, msg)
                        )

        except websockets.exceptions.ConnectionClosed:
            print("Studio disconnected.")

    # Register with ML-Dash server
    _register_bridge(url, api_key, port, token)

    # Start heartbeat task
    heartbeat_task = asyncio.create_task(_heartbeat_loop(url, api_key))

    # Start WebSocket server
    print(f"Connected to ML-Dash at {url}")
    print(f"Bridge running on ws://localhost:{port}")
    print(f"Code from ML-Dash Studio will execute on this machine.")
    print(f"Press Ctrl+C to stop.\n")

    # Use the stable websockets.serve API (works across versions)
    server = await websockets.serve(handler, "localhost", port)

    try:
        await asyncio.Future()  # Run forever
    except asyncio.CancelledError:
        pass
    finally:
        heartbeat_task.cancel()
        await kernel.stop()
        server.close()
        await server.wait_closed()


async def _heartbeat_loop(url, api_key):
    """Send heartbeat to ML-Dash every 15 seconds."""
    endpoint = url.rstrip('/') + '/api/studio/bridge/heartbeat/'
    headers = {'X-API-Key': api_key}

    while True:
        try:
            await asyncio.sleep(15)
            await asyncio.to_thread(
                lambda: requests.post(endpoint, headers=headers, timeout=5)
            )
        except asyncio.CancelledError:
            break
        except Exception:
            pass


def _register_bridge(url, api_key, port, token):
    """Register the bridge with the ML-Dash server."""
    endpoint = url.rstrip('/') + '/api/studio/bridge/register/'
    headers = {'X-API-Key': api_key}
    payload = {'port': port, 'token': token}

    try:
        resp = requests.post(endpoint, json=payload, headers=headers, timeout=10)
        if resp.status_code == 401:
            print("Error: Invalid API key. Run `mldash login` to re-authenticate.")
            sys.exit(1)
        resp.raise_for_status()
    except requests.ConnectionError:
        print(f"Error: Could not connect to ML-Dash at {url}")
        sys.exit(1)
    except requests.RequestException as e:
        print(f"Error registering bridge: {e}")
        sys.exit(1)
