"""
WebSocket Tunnel — Authenticated connection between Runtime and Neura Cloud.

The tunnel:
1. Connects to cloud via WebSocket (authenticated with device token)
2. Receives commands from cloud (execute_python, upload_file, etc.)
3. Dispatches to local DockerExecutor / StreamlitManager / OllamaProxy
4. Sends results back to cloud

Protocol: JSON messages with {type, id, payload}
"""
import asyncio
import json
import logging
from typing import Dict, Any, Optional, Callable

try:
    import websockets
    HAS_WEBSOCKETS = True
except ImportError:
    HAS_WEBSOCKETS = False

from .config import RuntimeConfig, get_device_token
from .docker_executor import DockerExecutor
from .streamlit_manager import StreamlitManager
from .ollama_proxy import OllamaProxy

logger = logging.getLogger(__name__)


class RuntimeTunnel:
    """WebSocket tunnel to Neura Cloud."""

    def __init__(
        self,
        config: RuntimeConfig,
        executor: DockerExecutor,
        streamlit: StreamlitManager,
        ollama: OllamaProxy,
    ):
        self.config = config
        self.executor = executor
        self.streamlit = streamlit
        self.ollama = ollama
        self._ws = None
        self._running = False
        self._reconnect_delay = 1  # exponential backoff

    async def connect(self):
        """Connect to cloud and start message loop."""
        if not HAS_WEBSOCKETS:
            logger.error("websockets package not installed. Run: pip install websockets")
            return

        token = get_device_token()
        if not token:
            logger.error("Device not paired. Run: neura-runtime pair")
            return

        self._running = True
        while self._running:
            try:
                ws_url = f"{self.config.cloud_ws_url}?token={token}"
                logger.info(f"Connecting to {self.config.cloud_ws_url}...")

                async with websockets.connect(ws_url, ping_interval=30) as ws:
                    self._ws = ws
                    self._reconnect_delay = 1  # reset on success
                    logger.info("Connected to Neura Cloud")

                    # Send capabilities
                    await self._send_capabilities()

                    # Message loop
                    async for message in ws:
                        try:
                            msg = json.loads(message)
                            asyncio.create_task(self._handle_message(msg))
                        except json.JSONDecodeError:
                            logger.warning(f"Invalid message: {message[:100]}")

            except Exception as e:
                logger.warning(f"Tunnel disconnected: {e}")
                if self._running:
                    logger.info(f"Reconnecting in {self._reconnect_delay}s...")
                    await asyncio.sleep(self._reconnect_delay)
                    self._reconnect_delay = min(self._reconnect_delay * 2, 30)

    async def disconnect(self):
        """Disconnect from cloud."""
        self._running = False
        if self._ws:
            await self._ws.close()
            self._ws = None

    async def _send_capabilities(self):
        """Report Runtime capabilities to cloud."""
        # Discover local models
        models = await self.ollama.discover()

        docker_ok = await self.executor.check_docker()

        await self._send({
            "type": "capabilities",
            "payload": {
                "version": "0.1.0",
                "docker_available": docker_ok,
                "local_models": models,
                "streamlit_ports": f"{self.config.streamlit_port_start}-{self.config.streamlit_port_end}",
                "streamlit_apps": self.streamlit.list_apps(),
            },
        })

    async def _handle_message(self, msg: Dict[str, Any]):
        """Dispatch incoming cloud commands."""
        msg_type = msg.get("type")
        msg_id = msg.get("id")
        payload = msg.get("payload", {})

        try:
            if msg_type == "execute_python":
                result = await self.executor.execute_python(
                    session_id=payload["session_id"],
                    code=payload["code"],
                    timeout=payload.get("timeout", 300),
                    packages=payload.get("packages"),
                )

            elif msg_type == "upload_file":
                import base64
                content = base64.b64decode(payload["content_b64"])
                result = await self.executor.upload_file(
                    session_id=payload["session_id"],
                    filename=payload["filename"],
                    content=content,
                    directory=payload.get("directory", "/home/sandbox"),
                )

            elif msg_type == "download_file":
                result = await self.executor.download_file(
                    session_id=payload["session_id"],
                    path=payload["path"],
                )
                # Encode binary content for WebSocket
                if result.get("success") and "content" in result:
                    import base64
                    result["content_b64"] = base64.b64encode(result.pop("content")).decode()

            elif msg_type == "create_container":
                result = await self.executor.create_container(
                    session_id=payload["session_id"],
                    template=payload.get("template", "default"),
                )

            elif msg_type == "destroy_container":
                ok = await self.executor.destroy_container(payload["session_id"])
                result = {"success": ok}

            elif msg_type == "launch_streamlit":
                result = await self.streamlit.launch(
                    app_id=payload["app_id"],
                    session_id=payload["session_id"],
                    app_code=payload["app_code"],
                    requirements=payload.get("requirements"),
                    title=payload.get("title", "Streamlit App"),
                )

            elif msg_type == "stop_streamlit":
                ok = await self.streamlit.stop(payload["app_id"])
                result = {"success": ok}

            elif msg_type == "list_streamlit":
                result = {"apps": self.streamlit.list_apps()}

            elif msg_type == "chat_completion":
                result = await self.ollama.chat_completion(
                    model_id=payload["model_id"],
                    messages=payload["messages"],
                    stream=payload.get("stream", False),
                )

            elif msg_type == "discover_models":
                models = await self.ollama.discover()
                result = {"models": models}

            elif msg_type == "ping":
                result = {"pong": True}

            else:
                result = {"error": f"Unknown command: {msg_type}"}

        except Exception as e:
            logger.error(f"Error handling {msg_type}: {e}")
            result = {"error": str(e)}

        # Send response
        await self._send({
            "type": f"{msg_type}_response",
            "id": msg_id,
            "payload": result,
        })

    async def _send(self, msg: Dict[str, Any]):
        """Send a message to cloud."""
        if self._ws:
            try:
                await self._ws.send(json.dumps(msg))
            except Exception as e:
                logger.warning(f"Failed to send message: {e}")
