"""OpenSandbox backend implementation using official SDK."""

import asyncio
import atexit
import os
import subprocess
import time
from datetime import timedelta
from pathlib import Path
from typing import Any

import httpx
from loguru import logger

from vikingbot.config.schema import SandboxConfig, SessionKey
from vikingbot.sandbox.backends import register_backend
from vikingbot.sandbox.base import SandboxBackend, SandboxNotStartedError

# Global to track the opensandbox-server process
_OSB_SERVER_PROCESS: "subprocess.Popen | None" = None


def _is_kubernetes_env() -> bool:
    if "KUBERNETES_SERVICE_HOST" in os.environ:
        return True
    if Path("/var/run/secrets/kubernetes.io/serviceaccount").exists():
        return True
    return False


async def _wait_for_server(url: str, timeout: int = 60) -> bool:
    logger.info("Waiting for OpenSandbox server at {}...", url)
    start_time = time.time()
    while time.time() - start_time < timeout:
        try:
            async with httpx.AsyncClient(timeout=5.0) as client:
                health_url = f"{url}/health"
                logger.debug("Checking health at: {}", health_url)
                response = await client.get(health_url)
                logger.debug("Health check response: {} - {}", response.status_code, response.text)
                if response.status_code == 200:
                    logger.info("OpenSandbox server is ready!")
                    return True
        except Exception as e:
            logger.debug("Health check failed: {}", e)
        await asyncio.sleep(1)
    logger.warning("OpenSandbox server not ready after {}s", timeout)
    return False


def _start_opensandbox_server() -> "subprocess.Popen | None":
    global _OSB_SERVER_PROCESS

    if _OSB_SERVER_PROCESS is not None:
        if _OSB_SERVER_PROCESS.poll() is None:
            logger.info("OpenSandbox server already running")
            return _OSB_SERVER_PROCESS
        else:
            logger.warning("OpenSandbox server process died, restarting")
            _OSB_SERVER_PROCESS = None

    try:
        config_path = Path.home() / ".sandbox.toml"
        if not config_path.exists():
            logger.info("Initializing OpenSandbox config at {}", config_path)
            try:
                result = subprocess.run(
                    ["opensandbox-server", "init-config", str(config_path), "--example", "docker"],
                    capture_output=True,
                    text=True,
                )
                if result.returncode != 0:
                    logger.warning("Failed to init config with --example, trying without...")
                    result = subprocess.run(
                        ["opensandbox-server", "init-config", str(config_path)],
                        capture_output=True,
                        text=True,
                    )
                    if result.returncode != 0:
                        logger.warning("Failed to init config, stderr: {}", result.stderr)
            except Exception as e:
                logger.warning("Failed to run init-config: {}", e)

        logger.info("Starting OpenSandbox server...")
        _OSB_SERVER_PROCESS = subprocess.Popen(
            ["opensandbox-server"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
        )

        return _OSB_SERVER_PROCESS
    except FileNotFoundError:
        logger.error("opensandbox-server command not found.")
        logger.error("Please start it manually first: opensandbox-server")
        return None
    except Exception as e:
        logger.error("Failed to start OpenSandbox server: {}", e)
        return None


def cleanup_opensandbox_server():
    global _OSB_SERVER_PROCESS
    if _OSB_SERVER_PROCESS is not None and _OSB_SERVER_PROCESS.poll() is None:
        logger.info("Stopping OpenSandbox server...")
        _OSB_SERVER_PROCESS.terminate()
        try:
            _OSB_SERVER_PROCESS.wait(timeout=5)
        except subprocess.TimeoutExpired:
            _OSB_SERVER_PROCESS.kill()
        _OSB_SERVER_PROCESS = None


# Register cleanup on exit
atexit.register(cleanup_opensandbox_server)


@register_backend("opensandbox")
class OpenSandboxBackend(SandboxBackend):
    def __init__(self, config: "SandboxConfig", session_key: SessionKey, workspace: Path):
        # OpenSandbox has built-in isolation, restrict_to_workspace is not needed
        super().__init__()
        self.config = config
        self.session_key = session_key
        self._workspace = workspace
        self._sandbox = None
        self._connection_config = None

        self._osb_config = config.backends.opensandbox

        self._is_vke = _is_kubernetes_env()
        if self._is_vke:
            self._server_url = "http://opensandbox-server:8080"
            logger.info(
                "Detected VKE environment, using OpenSandbox server at: {}", self._server_url
            )
        else:
            self._server_url = self._osb_config.server_url
            logger.info(
                "Detected local environment, using OpenSandbox server at: {}", self._server_url
            )

    async def start(self) -> None:
        self._workspace.mkdir(parents=True, exist_ok=True)

        if not self._is_vke:
            server_process = _start_opensandbox_server()
            if server_process:
                ready = await _wait_for_server(self._server_url, timeout=10)
                if not ready:
                    logger.info(
                        "OpenSandbox server not ready. Please start it manually: opensandbox-server"
                    )

        try:
            from opensandbox.config import ConnectionConfig
            from opensandbox.sandbox import Sandbox

            self._connection_config = ConnectionConfig(
                domain=self._server_url,
                api_key=self._osb_config.api_key,
                request_timeout=timedelta(seconds=300),
            )

            timeout_seconds = self._osb_config.runtime.timeout

            # Configure volumes
            volumes = None
            if not self._is_vke:
                # Local environment: mount host volume
                from opensandbox.models.sandboxes import Host, Volume

                volumes = [
                    Volume(
                        name="workspace",
                        host=Host(path=str(self._workspace.resolve())),
                        mountPath="/workspace",
                    )
                ]
            else:
                # VKE environment: always mount TOS PVC to /workspace
                try:
                    from opensandbox.models.sandboxes import Volume

                    # Build Volume with PVC using dictionary approach for compatibility
                    volume_dict = {
                        "name": "tos-workspace",
                        "persistentVolumeClaim": {"claimName": "vikingbot-data"},
                        "mountPath": "/workspace",
                    }

                    # Try to create Volume object
                    volumes = [Volume(**volume_dict)]
                    logger.info("Configured TOS PVC mount: vikingbot-data -> /workspace")
                except Exception as e:
                    logger.warning("Failed to create Volume object with PVC, falling back: {}", e)
                    volumes = None

            self._sandbox = await Sandbox.create(
                self._osb_config.default_image,
                connection_config=self._connection_config,
                timeout=timedelta(seconds=timeout_seconds),
                volumes=volumes,
            )

            logger.info("OpenSandbox created successfully")

        except ImportError:
            logger.error(
                "opensandbox SDK not installed. Install with: uv pip install 'openviking[bot-sandbox]' (or uv pip install -e \".[bot-sandbox]\" for local dev)"
            )
            raise
        except Exception as e:
            logger.error("Failed to create OpenSandbox: {}", e)
            import traceback

            logger.error("Full traceback:\n{}", traceback.format_exc())
            raise

    async def execute(self, command: str, timeout: int = 60, **kwargs: Any) -> str:
        if not self._sandbox:
            raise SandboxNotStartedError()

        logger.info("[OpenSandbox] Executing: {}", repr(command))

        if command.strip() == "pwd":
            return "/workspace" if self._is_vke else "/"

        try:
            from opensandbox.models.execd import RunCommandOpts

            opts = RunCommandOpts(timeout=timedelta(seconds=timeout))
            execution = await self._sandbox.commands.run(command, opts=opts)

            output_parts = []

            stdout_text = ""
            if execution.logs and execution.logs.stdout:
                stdout_text = "\n".join(
                    [chunk.text for chunk in execution.logs.stdout if chunk.text]
                )

            stderr_text = ""
            if execution.logs and execution.logs.stderr:
                stderr_text = "\n".join(
                    [chunk.text for chunk in execution.logs.stderr if chunk.text]
                )

            exit_code = execution.exit_code if hasattr(execution, "exit_code") else 0

            if stdout_text:
                output_parts.append(stdout_text)
            if stderr_text:
                output_parts.append(f"STDERR:\n{stderr_text}")
            if exit_code != 0:
                output_parts.append(f"\nExit code: {exit_code}")

            result = "\n".join(output_parts) if output_parts else "(no output)"

            max_len = 10000
            if len(result) > max_len:
                result = result[:max_len] + f"\n... (truncated, {len(result) - max_len} more chars)"

            logger.info("[OpenSandbox] Output:\n{}", result)
            return result

        except Exception as e:
            logger.error("[OpenSandbox] Error: {}", e)
            import traceback

            logger.error("[OpenSandbox] Traceback:\n{}", traceback.format_exc())
            raise

    async def stop(self) -> None:
        if self._sandbox:
            try:
                if hasattr(self._sandbox, "kill"):
                    await self._sandbox.kill()
                if hasattr(self._sandbox, "close"):
                    await self._sandbox.close()
                logger.info("OpenSandbox stopped")
            except Exception as e:
                logger.warning("Error stopping sandbox: {}", e)

        self._sandbox = None
        self._connection_config = None

    def is_running(self) -> bool:
        return self._sandbox is not None

    @property
    def workspace(self) -> Path:
        return self._workspace

    @property
    def sandbox_cwd(self) -> str:
        return "/workspace"

    async def read_file(self, path: str) -> str:
        """Read file from OpenSandbox."""
        if not self._sandbox:
            raise SandboxNotStartedError()

        # In VKE environment, use SDK API; in local, use base implementation (host mount)
        if self._is_vke:
            try:
                sandbox_path = path
                if not path.startswith("/"):
                    sandbox_path = f"/workspace/{path}"
                return await self._sandbox.files.read_file(sandbox_path)
            except Exception as e:
                logger.error(f"[OpenSandbox] Failed to read file {path}: {e}")
                raise
        else:
            return await super().read_file(path)

    async def write_file(self, path: str, content: str) -> None:
        """Write file to OpenSandbox."""
        if not self._sandbox:
            raise SandboxNotStartedError()

        # In VKE environment, use SDK API; in local, use base implementation (host mount)
        if self._is_vke:
            try:
                sandbox_path = path
                if not path.startswith("/"):
                    sandbox_path = f"/workspace/{path}"
                await self._sandbox.files.write_file(sandbox_path, content, mode=0o644)
            except Exception as e:
                logger.error(f"[OpenSandbox] Failed to write file {path}: {e}")
                raise
        else:
            await super().write_file(path, content)

    async def list_dir(self, path: str) -> list[tuple[str, bool]]:
        """List directory in OpenSandbox."""
        if not self._sandbox:
            raise SandboxNotStartedError()

        # In VKE environment, use SDK API; in local, use base implementation (host mount)
        if self._is_vke:
            try:
                sandbox_path = path
                if not path.startswith("/"):
                    sandbox_path = f"/workspace/{path}"

                # Use execute to list directory as fallback
                result = await self.execute(f"ls -la {sandbox_path}")

                items = []
                lines = result.strip().split("\n")
                for line in lines[1:]:
                    parts = line.split()
                    if len(parts) >= 9:
                        name = " ".join(parts[8:])
                        is_dir = parts[0].startswith("d")
                        if name not in (".", ".."):
                            items.append((name, is_dir))

                return items
            except Exception as e:
                logger.error(f"[OpenSandbox] Failed to list directory {path}: {e}")
                raise
        else:
            return await super().list_dir(path)
