from __future__ import annotations

import functools
import shlex
from typing import TYPE_CHECKING, Any

from shipyard import ShipyardClient, Spec

from astrbot.api import logger

if TYPE_CHECKING:
    from astrbot.core.agent.tool import ToolSchema

from astrbot.core.computer.olayer import (
    FileSystemComponent,
    PythonComponent,
    ShellComponent,
)

from .base import ComputerBooter
from .shell_background import build_detached_shell_command
from .shipyard_search_file_util import search_files_via_shell


def _maybe_model_dump(value: Any) -> dict[str, Any]:
    if isinstance(value, dict):
        return value
    if hasattr(value, "model_dump"):
        dumped = value.model_dump()
        if isinstance(dumped, dict):
            return dumped
    return {}


class ShipyardShellWrapper:
    def __init__(self, _shipyard_shell: ShellComponent):
        self._shell = _shipyard_shell

    async def exec(
        self,
        command: str,
        cwd: str | None = None,
        env: dict[str, str] | None = None,
        timeout: int | None = 300,
        shell: bool = True,
        background: bool = False,
        session_id: str | None = None,
    ) -> dict[str, Any]:
        _ = session_id
        if not shell:
            return {
                "stdout": "",
                "stderr": "error: only shell mode is supported in shipyard booter.",
                "exit_code": 2,
                "success": False,
            }

        run_command = command
        if env:
            env_prefix = " ".join(
                f"{k}={shlex.quote(str(v))}" for k, v in sorted(env.items())
            )
            run_command = f"{env_prefix} {run_command}"

        if background:
            run_command = build_detached_shell_command(run_command)

        result = await self._shell.exec(
            run_command,
            timeout=timeout or 300,
            cwd=cwd,
        )
        payload = _maybe_model_dump(result)

        stdout = payload.get("output", payload.get("stdout", "")) or ""
        stderr = payload.get("error", payload.get("stderr", "")) or ""
        exit_code = payload.get("exit_code")
        if background:
            pid: int | None = None
            try:
                pid = int(str(stdout).strip().splitlines()[-1])
            except Exception:
                pid = None
            return {
                "pid": pid,
                "stdout": (
                    f"Command is running in the background. pid={pid}"
                    if pid is not None
                    else "Command was submitted in the background."
                ),
                "stderr": stderr,
                "exit_code": exit_code,
                "success": bool(payload.get("success", not stderr)),
                "execution_id": payload.get("execution_id"),
                "execution_time_ms": payload.get("execution_time_ms"),
                "command": payload.get("command"),
            }

        return {
            "stdout": stdout,
            "stderr": stderr,
            "exit_code": exit_code,
            "success": bool(payload.get("success", not stderr)),
            "execution_id": payload.get("execution_id"),
            "execution_time_ms": payload.get("execution_time_ms"),
            "command": payload.get("command"),
        }


class ShipyardFileSystemWrapper:
    def __init__(
        self,
        _shipyard_fs: FileSystemComponent,
        _shipyard_shell: ShellComponent,
    ):
        self._fs = _shipyard_fs
        self._shell = _shipyard_shell

    async def create_file(
        self,
        path: str,
        content: str = "",
        mode: int = 420,
    ) -> dict[str, Any]:
        return await self._fs.create_file(path=path, content=content, mode=mode)

    async def read_file(
        self,
        path: str,
        encoding: str = "utf-8",
        offset: int | None = None,
        limit: int | None = None,
    ) -> dict[str, Any]:
        return await self._fs.read_file(
            path=path,
            encoding=encoding,
            offset=offset,
            limit=limit,
        )

    async def write_file(
        self,
        path: str,
        content: str,
        mode: str = "w",
        encoding: str = "utf-8",
    ) -> dict[str, Any]:
        return await self._fs.write_file(
            path=path,
            content=content,
            mode=mode,
            encoding=encoding,
        )

    async def list_dir(
        self,
        path: str = ".",
        show_hidden: bool = False,
    ) -> dict[str, Any]:
        return await self._fs.list_dir(path=path, show_hidden=show_hidden)

    async def delete_file(self, path: str) -> dict[str, Any]:
        return await self._fs.delete_file(path=path)

    async def search_files(
        self,
        pattern: str,
        path: str | None = None,
        glob: str | None = None,
        after_context: int | None = None,
        before_context: int | None = None,
    ) -> dict[str, Any]:
        return await search_files_via_shell(
            self._shell,
            pattern=pattern,
            path=path,
            glob=glob,
            after_context=after_context,
            before_context=before_context,
        )

    async def edit_file(
        self,
        path: str,
        old_string: str,
        new_string: str,
        replace_all: bool = False,
        encoding: str = "utf-8",
    ) -> dict[str, Any]:
        return await self._fs.edit_file(
            path=path,
            old_string=old_string,
            new_string=new_string,
            replace_all=replace_all,
            encoding=encoding,
        )


class ShipyardBooter(ComputerBooter):
    @classmethod
    @functools.cache
    def _default_tools(cls) -> tuple[ToolSchema, ...]:
        from astrbot.core.computer.tools import (
            ExecuteShellTool,
            FileDownloadTool,
            FileUploadTool,
            PythonTool,
        )

        return (
            ExecuteShellTool(),
            PythonTool(),
            FileUploadTool(),
            FileDownloadTool(),
        )

    @classmethod
    def get_default_tools(cls) -> list[ToolSchema]:
        return list(cls._default_tools())

    def __init__(
        self,
        endpoint_url: str,
        access_token: str,
        ttl: int = 3600,
        session_num: int = 10,
    ) -> None:
        self._sandbox_client = ShipyardClient(
            endpoint_url=endpoint_url,
            access_token=access_token,
        )
        self._ttl = ttl
        self._session_num = session_num

    async def boot(self, session_id: str) -> None:
        ship = await self._sandbox_client.create_ship(
            ttl=self._ttl,
            spec=Spec(cpus=1.0, memory="512m"),
            max_session_num=self._session_num,
            session_id=session_id,
        )
        logger.info(
            "[Computer] sandbox_created booter=shipyard ship_id=%s session=%s",
            ship.id,
            session_id,
        )
        self._ship = ship
        self._shell = ShipyardShellWrapper(self._ship.shell)  # type: ignore[arg-type]
        self._fs = ShipyardFileSystemWrapper(self._ship.fs, self._shell)  # type: ignore[arg-type]

    async def shutdown(self, **kwargs) -> None:
        logger.info("[Computer] booter_shutdown booter=shipyard status=done")

    @property
    def fs(self) -> FileSystemComponent:
        return self._ship.fs  # type: ignore[return-value]

    @property
    def python(self) -> PythonComponent:
        return self._ship.python

    @property
    def shell(self) -> ShellComponent:
        return self._shell

    async def upload_file(self, path: str, file_name: str) -> dict:
        """Upload file to sandbox"""
        result = await self._ship.upload_file(path, file_name)
        logger.info(
            "[Computer] file_upload booter=shipyard remote_path=%s",
            file_name,
        )
        return result

    async def download_file(self, remote_path: str, local_path: str):
        """Download file from sandbox."""
        result = await self._ship.download_file(remote_path, local_path)
        logger.info(
            "[Computer] file_download booter=shipyard remote_path=%s local_path=%s",
            remote_path,
            local_path,
        )
        return result

    async def available(self) -> bool:
        """Check if the sandbox is available."""
        try:
            ship_id = self._ship.id
            data = await self._sandbox_client.get_ship(ship_id)
            if not data:
                logger.debug(
                    "[Computer] health_check booter=shipyard ship_id=%s healthy=false reason=no_data",
                    ship_id,
                )
                return False
            health = bool(data.get("status", 0) == 1)
            logger.debug(
                "[Computer] health_check booter=shipyard ship_id=%s healthy=%s",
                ship_id,
                health,
            )
            return health
        except Exception:
            logger.exception(
                "[Computer] health_check_failed booter=shipyard ship_id=%s",
                getattr(getattr(self, "_ship", None), "id", "unknown"),
            )
            return False