"""
CommandsManager - 通过 envd 兼容层（port 49983）在 sandbox 内执行命令

envd 使用 Connect-RPC 协议：
  - POST /process.Process/Start  → 启动进程，返回 {"pid": "..."}
  - POST /process.Process/Connect → 连接进程，流式读取 stdout/stderr/exit
  - POST /process.Process/Kill   → 终止进程

帧格式（Connect-protocol）：
  [1 byte flags][4 bytes big-endian body length][JSON body]

输出读取策略（解决 Connect 404 问题）：
  envd 进程结束后立即从 store 移除，导致 Connect 返回 404 且输出丢失。
  解决方案：同步命令将 stdout/stderr 重定向到临时文件，通过 /files REST 接口读取；
  Connect 仅用作"等待进程退出"的信号（404 即表示已退出，继续读文件即可）。
"""

import struct
from dataclasses import dataclass
from shlex import quote as _shell_quote
from typing import TYPE_CHECKING, Optional
from urllib.parse import quote as _url_quote

import httpx

if TYPE_CHECKING:
    from .sandbox import Sandbox


@dataclass
class CommandResult:
    """同步命令执行结果"""
    stdout: str = ""
    stderr: str = ""
    exit_code: int = 0
    pid: str = ""


class CommandExitException(Exception):
    """命令以非零退出码退出时抛出（可选使用）"""

    def __init__(self, exit_code: int, stdout: str, stderr: str):
        self.exit_code = exit_code
        self.stdout = stdout
        self.stderr = stderr
        super().__init__(f"Command exited with code {exit_code}")


class CommandHandle:
    """后台进程 handle，由 commands.run(..., background=True) 返回"""

    def __init__(self, pid: str, sandbox: "Sandbox"):
        self.pid = pid
        self._sandbox = sandbox

    def kill(self):
        """终止后台进程"""
        try:
            headers = self._sandbox._get_envd_headers()
            base = self._sandbox._get_envd_base_url()
            httpx.post(
                f"{base}/process.Process/Kill",
                json={"process": {"pid": self.pid}},
                headers=headers,
                timeout=10,
            )
        except Exception:
            pass


# ---------------------------------------------------------------------------
# 内部辅助函数
# ---------------------------------------------------------------------------

def _parse_pid(raw: bytes) -> Optional[str]:
    """从 Start 响应中解析 pid（支持 Connect-RPC 帧和裸 JSON 两种格式）"""
    import json
    pid = None
    # 优先按 Connect-RPC 帧解析
    if len(raw) >= 5:
        body_len = struct.unpack(">I", raw[1:5])[0]
        try:
            frame = json.loads(raw[5:5 + body_len])
            pid = frame.get("pid") or frame.get("process", {}).get("pid")
        except Exception:
            pass
    # 降级：裸 JSON
    if not pid:
        try:
            frame = json.loads(raw)
            pid = frame.get("pid")
        except Exception:
            pass
    return pid


def _wait_for_exit(base_url: str, pid: str, headers: dict, timeout: float) -> None:
    """
    通过 Connect 等待进程退出。
      - 收到 exit 事件（result.exitCode）→ 进程已结束，返回
      - HTTP 404 → 进程已从 store 移除（已退出）→ 直接返回
      - 超时 → 静默返回（后续仍可读文件）
    """
    import json
    try:
        with httpx.stream(
            "POST",
            f"{base_url}/process.Process/Connect",
            json={"process": {"pid": pid}},
            headers=headers,
            timeout=timeout,
        ) as resp:
            if resp.status_code == 404:
                return  # 进程已退出
            if resp.status_code >= 300:
                return
            buf = b""
            for raw_chunk in resp.iter_raw():
                buf += raw_chunk
                while len(buf) >= 5:
                    body_len = struct.unpack(">I", buf[1:5])[0]
                    frame_size = 5 + body_len
                    if len(buf) < frame_size:
                        break
                    frame_bytes = buf[5:frame_size]
                    buf = buf[frame_size:]
                    try:
                        frame = json.loads(frame_bytes)
                    except Exception:
                        continue
                    if "result" in frame:  # exit event
                        return
    except httpx.TimeoutException:
        pass
    except Exception:
        pass


def _read_output_files(
    base_url: str,
    headers: dict,
    stdout_file: str,
    stderr_file: str,
    exit_file: str,
    pid: str,
) -> "CommandResult":
    """通过 /files REST 接口读取命令的 stdout/stderr/exit_code 临时文件"""

    def _read_file(path: str) -> str:
        try:
            resp = httpx.get(
                f"{base_url}/files?path={_url_quote(path)}",
                headers=headers,
                timeout=15,
            )
            if resp.status_code == 200:
                return resp.text
        except Exception:
            pass
        return ""

    result = CommandResult(pid=pid)
    result.stdout = _read_file(stdout_file)
    result.stderr = _read_file(stderr_file)
    exit_str = _read_file(exit_file).strip()
    try:
        result.exit_code = int(exit_str)
    except (ValueError, TypeError):
        result.exit_code = 0

    # 清理临时文件（异步 Start 触发 rm，忽略失败）
    for path in (stdout_file, stderr_file, exit_file):
        try:
            httpx.post(
                f"{base_url}/process.Process/Start",
                json={"cmd": {"args": ["rm", "-f", path]},
                      "tag": "sdk-cleanup"},
                headers=headers,
                timeout=5,
            )
        except Exception:
            pass

    return result


# ---------------------------------------------------------------------------
# CommandsManager
# ---------------------------------------------------------------------------

class CommandsManager:
    """在 sandbox 中执行 shell 命令（通过 envd 兼容层）"""

    def __init__(self, sandbox: "Sandbox"):
        self._sandbox = sandbox

    def run(
        self,
        cmd: str,
        background: bool = False,
        timeout: int = 60,
    ):
        """
        执行 shell 命令。

        策略（同步模式）：
          1. 包装命令：stdout/stderr 重定向到 /tmp 临时文件，退出码写入另一个文件
          2. Start 启动包装命令
          3. Connect 等待退出（404 = 已退出，同样视为完成）
          4. 通过 /files REST 读取三个临时文件，组装 CommandResult

        :param cmd: shell 命令字符串
        :param background: 若为 True，在后台运行并立即返回 CommandHandle
        :param timeout: 超时秒数（仅同步模式有效）
        :return: CommandResult（同步）或 CommandHandle（后台）
        """
        import uuid as _uuid
        base = self._sandbox._get_envd_base_url()
        headers = self._sandbox._get_envd_headers()

        # ---- 后台模式：直接 Start，不关心输出 ----
        if background:
            start_resp = httpx.post(
                f"{base}/process.Process/Start",
                json={"cmd": {"args": ["sh", "-c", cmd]}, "tag": "sdk"},
                headers=headers,
                timeout=30,
            )
            start_resp.raise_for_status()
            pid = _parse_pid(start_resp.content)
            if not pid:
                raise RuntimeError(
                    f"Failed to start background process, response: {start_resp.content[:200]}"
                )
            return CommandHandle(pid=pid, sandbox=self._sandbox)

        # ---- 同步模式：输出重定向到临时文件 ----
        run_id = str(_uuid.uuid4())[:8]
        stdout_file = f"/tmp/.sdk_stdout_{run_id}"
        stderr_file = f"/tmp/.sdk_stderr_{run_id}"
        exit_file = f"/tmp/.sdk_exit_{run_id}"

        # 包装命令：先执行原始命令，将 stdout/stderr 重定向，最后写 exit code
        wrapped = (
            f"sh -c {_shell_quote(cmd)} "
            f">{stdout_file} 2>{stderr_file}; "
            f"echo $? >{exit_file}"
        )

        start_resp = httpx.post(
            f"{base}/process.Process/Start",
            json={"cmd": {"args": ["sh", "-c", wrapped]}, "tag": "sdk"},
            headers=headers,
            timeout=30,
        )
        start_resp.raise_for_status()

        pid = _parse_pid(start_resp.content)
        if not pid:
            raise RuntimeError(
                f"Failed to start process, response: {start_resp.content[:200]}"
            )

        # 等待进程退出（Connect 404 = 已退出，直接读文件）
        _wait_for_exit(base, pid, headers, float(timeout))

        # 通过 /files 读取输出
        return _read_output_files(
            base_url=base,
            headers=headers,
            stdout_file=stdout_file,
            stderr_file=stderr_file,
            exit_file=exit_file,
            pid=pid,
        )
