"""WebSocket session for AIR one-shot / deep-research execution.

Connects to the backend via WebSocket, submits a task, and handles the
full message loop — executing code locally, installing packages, and
writing files until the backend signals completion.

This is a Python port of air-ui's ``useWebSocket.ts``.
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass, field
from typing import Any, Callable
from urllib.parse import urlencode

import websockets
import websockets.asyncio.client

from .executor import ExecutionResult, FileInfo, LocalCodeExecutor

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Result
# ---------------------------------------------------------------------------


@dataclass
class OneShotResult:
    """Returned by :meth:`AIR.one_shot` when the task finishes."""

    task_id: str
    output: str
    result: Any = None
    files_created: list[FileInfo] = field(default_factory=list)
    work_dir: str = ""
    error: str | None = None


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _extract_task_id(backend_path: str) -> str | None:
    """Extract a ``task_*`` segment from a backend work_dir path."""
    for part in reversed(backend_path.split("/")):
        if part.startswith("task_"):
            return part
    return None


def _http_to_ws(url: str) -> str:
    """Convert an ``http(s)://`` URL to ``ws(s)://``."""
    if url.startswith("https://"):
        return "wss://" + url[len("https://"):]
    if url.startswith("http://"):
        return "ws://" + url[len("http://"):]
    return url


# ---------------------------------------------------------------------------
# WebSocketSession
# ---------------------------------------------------------------------------


class WebSocketSession:
    """Manage a single WebSocket task session with the AIR backend.

    Parameters
    ----------
    base_url:
        HTTP(S) base URL of the backend (e.g. ``http://localhost:8000``).
    api_key:
        AIR API key used for authentication.
    task_id:
        Unique identifier for this task.
    work_dir:
        Local directory where code and outputs are stored.
    on_output:
        Optional callback invoked with streaming text from the backend
        and from local code execution.
    timeout:
        Maximum wall-clock seconds for the entire session.
    python_path:
        Explicit Python binary for the venv.  ``None`` → auto-detect.
    venv_path:
        Explicit path to an existing virtual environment.  ``None`` →
        walk up from *work_dir* to find one, or create a new one.
    """

    def __init__(
        self,
        base_url: str,
        api_key: str,
        task_id: str,
        work_dir: str,
        on_output: Callable[[str], None] | None = None,
        timeout: int = 86400,
        python_path: str | None = None,
        venv_path: str | None = None,
        config: dict | None = None,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._task_id = task_id
        self._work_dir = work_dir
        self._on_output = on_output
        self._timeout = timeout
        self._python_path = python_path
        self._venv_path = venv_path
        self._config = config or {}

        self._executor: LocalCodeExecutor | None = None
        self._all_files: list[FileInfo] = []

    # ------------------------------------------------------------------
    # Public
    # ------------------------------------------------------------------

    async def run(self, task: str, config: dict) -> OneShotResult:
        """Connect, submit *task*, and drive the execution loop to completion."""
        # Build WS URL
        ws_base = _http_to_ws(self._base_url)
        params = urlencode({"api_key": self._api_key})
        ws_url = f"{ws_base}/ws/{self._task_id}?{params}"

        # Create executor
        self._executor = LocalCodeExecutor(
            self._work_dir,
            python_path=self._python_path,
            venv_path=self._venv_path,
        )
        await self._executor.initialize()

        output_parts: list[str] = []
        result_data: Any = None
        errors: list[str] = []

        def _out(text: str) -> None:
            output_parts.append(text)
            if self._on_output:
                self._on_output(text)

        try:
            async with websockets.asyncio.client.connect(
                ws_url,
                open_timeout=30,
                close_timeout=10,
                ping_interval=300,
                ping_timeout=3600,
                max_size=2**24,  # 16 MiB
            ) as ws:
                # Submit task
                await ws.send(json.dumps({
                    "type": "task_submit",
                    "task": task,
                    "config": config,
                }))

                # Message loop
                while True:
                    raw = await ws.recv()
                    msg = json.loads(raw)
                    msg_type = msg.get("type", "")

                    if msg_type == "output":
                        data = msg.get("data", "")
                        if data:
                            _out(data)

                    elif msg_type == "status":
                        text = msg.get("message", "")
                        if text:
                            _out(f"Status: {text}\n")

                    elif msg_type == "result":
                        result_data = msg.get("data")

                    elif msg_type == "error":
                        err = msg.get("message", "Unknown error")
                        errors.append(err)
                        _out(f"Error: {err}\n")

                    elif msg_type == "complete":
                        _out("Task execution completed\n")
                        break

                    elif msg_type == "execute_code":
                        await self._handle_execute_code(ws, msg, _out)

                    elif msg_type == "install_packages":
                        await self._handle_install_packages(ws, msg, _out)

                    elif msg_type == "write_file":
                        await self._handle_write_file(msg, _out)

                    elif msg_type == "heartbeat":
                        pass  # keep-alive

                    else:
                        logger.debug("Unknown message type: %s", msg_type)

        except websockets.exceptions.ConnectionClosed as exc:
            errors.append(f"WebSocket closed: {exc}")
        except Exception as exc:
            errors.append(str(exc))

        return OneShotResult(
            task_id=self._task_id,
            output="".join(output_parts),
            result=result_data,
            files_created=self._all_files,
            work_dir=self._work_dir,
            error="; ".join(errors) if errors else None,
        )

    # ------------------------------------------------------------------
    # Message handlers
    # ------------------------------------------------------------------

    async def _handle_execute_code(self, ws, msg: dict, out: Callable[[str], None]) -> None:
        execution_id = msg.get("execution_id", "")
        code_blocks = msg.get("code_blocks", [])
        timeout = msg.get("timeout", 86400)
        msg_work_dir = msg.get("work_dir", "")

        # Acknowledge immediately
        await ws.send(json.dumps({
            "type": "execution_ack",
            "execution_id": execution_id,
        }))

        n = len(code_blocks)
        out(f"\n--- Executing code ({n} block(s)) ---\n")

        assert self._executor is not None

        # Update executor work_dir to match the backend's phase subdirectory
        # (e.g. ".../task_xxx/control/" → run code in local task_dir/control/)
        effective_work_dir = self._work_dir
        if msg_work_dir:
            backend_task_id = _extract_task_id(msg_work_dir)
            if backend_task_id:
                idx = msg_work_dir.find(backend_task_id)
                suffix = msg_work_dir[idx + len(backend_task_id):].strip("/")
                if suffix:
                    effective_work_dir = os.path.join(self._work_dir, suffix)
        self._executor.work_dir = effective_work_dir
        for sub in ("codebase", "data", "reports"):
            os.makedirs(os.path.join(effective_work_dir, sub), exist_ok=True)

        try:
            result: ExecutionResult = await self._executor.execute_code_blocks(
                code_blocks,
                timeout=timeout,
                on_output=out,
            )

            # Sync files to backend BEFORE sending execution result,
            # so that post_execution_transfer can find images for VLM review
            if result.files_created:
                self._all_files.extend(result.files_created)
                await self._sync_files_to_backend(ws, result, msg_work_dir, effective_work_dir)
                await ws.send(json.dumps({
                    "type": "files_created",
                    "execution_id": execution_id,
                    "files": [
                        {"path": f.path, "size": f.size, "mime": f.mime}
                        for f in result.files_created
                    ],
                }))
                out(f"\nCreated {len(result.files_created)} file(s)\n")

            # Send result to backend (triggers post_execution_transfer)
            await ws.send(json.dumps({
                "type": "execution_result",
                "execution_id": execution_id,
                "result": {
                    "exit_code": result.exit_code,
                    "output": result.output,
                    "code_file": result.code_file,
                },
            }))

            out(f"\n--- Execution complete (exit code: {result.exit_code}) ---\n")

        except Exception as exc:
            error_msg = str(exc)
            out(f"\nExecution error: {error_msg}\n")
            await ws.send(json.dumps({
                "type": "execution_error",
                "execution_id": execution_id,
                "error": error_msg,
            }))

    async def _handle_install_packages(self, ws, msg: dict, out: Callable[[str], None]) -> None:
        packages = msg.get("packages", [])
        out(f"\n--- Installing packages: {', '.join(packages)} ---\n")

        assert self._executor is not None

        result = await self._executor.install_packages(packages, on_output=out)

        await ws.send(json.dumps({
            "type": "install_complete",
            "packages": packages,
            "success": result.success,
            "failed": result.failed,
        }))

        if result.success:
            out("Packages installed successfully\n")
        else:
            out(f"Failed packages: {', '.join(result.failed)}\n")

    async def _handle_write_file(self, msg: dict, out: Callable[[str], None]) -> None:
        file_path = msg.get("path", "")
        content = msg.get("content", "")
        encoding = msg.get("encoding", "utf-8")
        msg_work_dir = msg.get("work_dir", "")

        if not file_path:
            return

        # Determine the correct local subdirectory from the message's work_dir.
        # E.g. if work_dir=".../task_xxx/control/", prepend "control/" to path.
        # Write directly to the task root (self._work_dir) to avoid conflicts
        # with the executor's work_dir which may point to a phase subdirectory.
        prefix = ""
        if msg_work_dir:
            backend_task_id = _extract_task_id(msg_work_dir)
            if backend_task_id:
                idx = msg_work_dir.find(backend_task_id)
                suffix = msg_work_dir[idx + len(backend_task_id):].strip("/")
                if suffix:
                    prefix = suffix

        local_path = os.path.join(self._work_dir, prefix, file_path) if prefix else os.path.join(self._work_dir, file_path)
        display_path = f"{prefix}/{file_path}" if prefix else file_path

        try:
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            if encoding == "base64":
                import base64
                with open(local_path, "wb") as f:
                    f.write(base64.b64decode(content))
            else:
                with open(local_path, "w", encoding=encoding) as f:
                    f.write(content)
            out(f"File written: {display_path}\n")
        except Exception as exc:
            out(f"Failed to write file {display_path}: {exc}\n")

    async def _sync_files_to_backend(
        self, ws, result: ExecutionResult, msg_work_dir: str, effective_work_dir: str,
    ) -> None:
        """Send code/report content and data stubs back to the backend.

        After each execution round the backend's work_dir needs matching
        files so that ``record_status()`` (codebase info, workspace manifest,
        step code/reports) sees them on disk.

        * ``codebase/`` and ``reports/`` files → actual content
        * everything else → metadata-only stub (path + size)
        """
        import base64

        # Check if VLM review is enabled — if so, sync images with content
        vlm_enabled = self._config.get("enableVlmReview", False)
        _IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".gif", ".svg", ".pdf", ".webp"}

        sync_entries: list[dict] = []
        for f in result.files_created:
            full_path = os.path.join(effective_work_dir, f.path)
            if f.path.startswith(("codebase/", "reports/")) and os.path.isfile(full_path):
                try:
                    with open(full_path, "r", encoding="utf-8") as fh:
                        sync_entries.append({"path": f.path, "size": f.size, "content": fh.read()})
                except Exception:
                    sync_entries.append({"path": f.path, "size": f.size})
            elif vlm_enabled and os.path.splitext(f.path)[1].lower() in _IMAGE_EXTS and os.path.isfile(full_path):
                # Upload image content for server-side VLM review
                try:
                    with open(full_path, "rb") as fh:
                        b64 = base64.b64encode(fh.read()).decode("ascii")
                    sync_entries.append({"path": f.path, "size": f.size, "content_base64": b64})
                except Exception:
                    sync_entries.append({"path": f.path, "size": f.size})
            else:
                sync_entries.append({"path": f.path, "size": f.size})

        if sync_entries:
            await ws.send(json.dumps({
                "type": "sync_files",
                "work_dir": msg_work_dir,
                "files": sync_entries,
            }))
