"""沙箱客户端：封装对 Rust 沙箱服务的调用，含司命授权签发与一炁事件回写。

使用约束（见 backend/CLAUDE.md）：禁止在调用处裸 ``async with httpx.AsyncClient()``，
应复用注入的 ``AsyncClient``（长生命周期，由上层依赖注入管理）。
"""

from __future__ import annotations

import json
import time
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import Any, NoReturn, TypeVar

import httpx
from pydantic import BaseModel

from app.core.config import SandboxSettings
from app.integrations.sandbox.auth import AuthorizationContext, sign_authorization
from app.integrations.sandbox.events import (
    EventSink,
    LoggingEventSink,
    completed_event,
    failed_event,
    started_event,
)
from app.integrations.sandbox.exceptions import (
    SandboxBadRequest,
    SandboxError,
    SandboxExecFailed,
    SandboxNotFound,
    SandboxUnauthorized,
    SandboxUnavailable,
)
from app.integrations.sandbox.models import (
    CancelResponse,
    ExecRequest,
    ExecResponse,
    FileListRequest,
    FileListResponse,
    FileReadRequest,
    FileReadResponse,
    FileStatRequest,
    FileStatResponse,
    FileWriteRequest,
    FileWriteResponse,
)

_T = TypeVar("_T", bound=BaseModel)


@dataclass(frozen=True)
class SandboxStreamEvent:
    """沙箱流式事件：stdout/stderr 增量（base64 str，保真透传）或 done（结束汇总 dict）。"""

    kind: str  # "stdout" | "stderr" | "done"
    payload: str | dict[str, Any]


class SandboxClient:
    """沙箱服务异步客户端。

    授权语义：调用方需先经司命审核（在本类之外完成），通过 ``AuthorizationContext`` 携带
    授权上下文；本类据此签发短期 HMAC 凭证（exec 含本次请求的 body_hash 绑定），放入
    ``Authorization`` 头，由沙箱端校验。
    """

    def __init__(
        self,
        http_client: httpx.AsyncClient,
        settings: SandboxSettings,
        event_sink: EventSink | None = None,
    ) -> None:
        if settings.sandbox_auth_enabled and not settings.sandbox_auth_secret:
            raise ValueError("auth_enabled=True 时 auth_secret 不能为空（须与沙箱端共享）")
        self._http = http_client
        self._settings = settings
        self._sink: EventSink = event_sink or LoggingEventSink()

    async def health(self) -> bool:
        """健康探测，失败返回 False 而不抛异常。"""
        try:
            resp = await self._http.get(f"{self._settings.rust_sandbox_url}/healthz", timeout=5.0)
        except httpx.HTTPError:
            return False
        return resp.status_code == 200

    async def exec(
        self,
        request: ExecRequest,
        auth_ctx: AuthorizationContext,
        *,
        event_sink: EventSink | None = None,
    ) -> ExecResponse:
        """同步执行：收集全部输出后返回。授权失败、网络错误或沙箱错误分别抛出对应异常，
        并回写 tool.failed 事件。``event_sink`` 可按调用覆盖，便于编排层按 run 绑定上下文。"""
        sink = event_sink or self._sink
        trace_id = request.trace_id or auth_ctx.trace_id
        run_id = auth_ctx.run_id
        tool = auth_ctx.tool

        # 请求体带上解析后的 trace_id，便于沙箱日志关联
        if request.trace_id is None and trace_id is not None:
            request.trace_id = trace_id

        await sink.emit(
            started_event(
                tool=tool,
                argv=request.argv or [f"code:{request.language}", request.code],
                risk_level=auth_ctx.risk_level,
                trace_id=trace_id,
                run_id=run_id,
            )
        )

        try:
            resp = await self._http.post(
                f"{self._settings.rust_sandbox_url}/v1/exec",
                json=request.model_dump(exclude_none=True),
                headers=self._exec_auth_headers(auth_ctx, request),
                timeout=self._settings.http_timeout(request.timeout_ms),
            )
        except httpx.HTTPError as exc:
            await sink.emit(
                failed_event(
                    tool=tool,
                    trace_id=trace_id,
                    run_id=run_id,
                    error_code="sandbox_unavailable",
                    error_message=str(exc),
                )
            )
            raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc

        if resp.status_code >= 400:
            exc_type: type[SandboxError]
            if resp.status_code == 401:
                exc_type = SandboxUnauthorized
            elif resp.status_code == 400:
                exc_type = SandboxBadRequest
            else:
                exc_type = SandboxExecFailed
            await self._fail(tool, trace_id, run_id, resp, exc_type, sink)

        result = ExecResponse.model_validate(resp.json())
        await sink.emit(
            completed_event(
                tool=tool,
                trace_id=trace_id,
                run_id=run_id,
                exit_code=result.exit_code,
                duration_ms=result.duration_ms,
                timeout=result.timeout,
                truncated=result.truncated,
                stdout_len=len(result.stdout),
                stderr_len=len(result.stderr),
            )
        )
        return result

    async def exec_stream(
        self,
        request: ExecRequest,
        auth_ctx: AuthorizationContext,
    ) -> AsyncIterator[SandboxStreamEvent]:
        """流式执行：实时产出 stdout/stderr 增量（base64 str）与 done 汇总（dict）。

        消费沙箱 ``POST /v1/exec/stream`` 的 SSE。不回写一炁事件——started/completed/failed
        由编排层 ``ToolRuntime.execute_stream`` 负责，本方法只透传执行级输出。
        """
        if request.trace_id is None and auth_ctx.trace_id is not None:
            request.trace_id = auth_ctx.trace_id

        try:
            async with self._http.stream(
                "POST",
                f"{self._settings.rust_sandbox_url}/v1/exec/stream",
                json=request.model_dump(exclude_none=True),
                headers=self._exec_auth_headers(auth_ctx, request),
                timeout=self._settings.http_timeout(request.timeout_ms),
            ) as resp:
                if resp.status_code == 401:
                    raise SandboxUnauthorized(await self._read_error(resp))
                if resp.status_code == 400:
                    raise SandboxBadRequest(await self._read_error(resp))
                if resp.status_code >= 400:
                    raise SandboxExecFailed(await self._read_error(resp))
                async for event_type, data in self._iter_sse(resp):
                    if event_type in ("stdout", "stderr"):
                        yield SandboxStreamEvent(event_type, data)
                    elif event_type == "done":
                        try:
                            yield SandboxStreamEvent("done", json.loads(data))
                        except ValueError:
                            continue
        except httpx.HTTPError as exc:
            raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc

    async def cancel(
        self,
        trace_id: str,
        auth_ctx: AuthorizationContext,
    ) -> bool:
        """取消运行中的任务（按 trace_id）。命中返回 True；任务不存在或已完成抛 SandboxNotFound。

        cancel 不回写一炁事件（由编排层按需记录），凭证不绑定 body_hash（危害有限）。
        """
        try:
            resp = await self._http.post(
                f"{self._settings.rust_sandbox_url}/v1/cancel",
                json={"trace_id": trace_id},
                headers=self._auth_headers(auth_ctx),
                timeout=10.0,
            )
        except httpx.HTTPError as exc:
            raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc

        if resp.status_code == 404:
            raise SandboxNotFound(f"任务不存在或已完成: {trace_id}")
        if resp.status_code >= 400:
            code, message = self._extract_error(resp)
            if resp.status_code == 401:
                raise SandboxUnauthorized(message)
            raise SandboxExecFailed(message, code=code)
        return CancelResponse.model_validate(resp.json()).cancelled

    # ---- 文件 IO（不经 shell，二进制安全；与 Rust fileio.rs 契约对齐） ----

    async def file_read(
        self,
        request: FileReadRequest,
        auth_ctx: AuthorizationContext,
        *,
        event_sink: EventSink | None = None,
    ) -> FileReadResponse:
        return await self._file_op(
            "read", "/v1/file/read", request, FileReadResponse, auth_ctx, event_sink=event_sink
        )

    async def file_write(
        self,
        request: FileWriteRequest,
        auth_ctx: AuthorizationContext,
        *,
        event_sink: EventSink | None = None,
    ) -> FileWriteResponse:
        return await self._file_op(
            "write", "/v1/file/write", request, FileWriteResponse, auth_ctx, event_sink=event_sink
        )

    async def file_list(
        self,
        request: FileListRequest,
        auth_ctx: AuthorizationContext,
        *,
        event_sink: EventSink | None = None,
    ) -> FileListResponse:
        return await self._file_op(
            "list", "/v1/file/list", request, FileListResponse, auth_ctx, event_sink=event_sink
        )

    async def file_stat(
        self,
        request: FileStatRequest,
        auth_ctx: AuthorizationContext,
        *,
        event_sink: EventSink | None = None,
    ) -> FileStatResponse:
        return await self._file_op(
            "stat", "/v1/file/stat", request, FileStatResponse, auth_ctx, event_sink=event_sink
        )

    async def _file_op(
        self,
        op: str,
        endpoint: str,
        request: FileReadRequest | FileWriteRequest | FileListRequest | FileStatRequest,
        response_cls: type[_T],
        auth_ctx: AuthorizationContext,
        *,
        event_sink: EventSink | None = None,
    ) -> _T:
        """文件操作通用流程：签发 file_op 凭证 → POST → 解析响应，并回写一炁事件。

        与 exec 共用错误映射（``_fail``）与事件构造；argv 记 ``[file:<op>, <path>]``
        供审计 input_summary。文件操作在主进程运行（非 exec 子进程），无超时参数，
        用默认 HTTP 超时（``http_timeout(None)``）。
        """
        sink = event_sink or self._sink
        trace_id = auth_ctx.trace_id
        run_id = auth_ctx.run_id
        tool = auth_ctx.tool
        path = request.path

        await sink.emit(
            started_event(
                tool=tool,
                argv=[f"file:{op}", path],
                risk_level=auth_ctx.risk_level,
                trace_id=trace_id,
                run_id=run_id,
            )
        )

        t0 = time.monotonic()
        try:
            resp = await self._http.post(
                f"{self._settings.rust_sandbox_url}{endpoint}",
                json=request.model_dump(exclude_none=True),
                headers=self._file_auth_headers(auth_ctx, op, path),
                timeout=self._settings.http_timeout(None),
            )
        except httpx.HTTPError as exc:
            await sink.emit(
                failed_event(
                    tool=tool,
                    trace_id=trace_id,
                    run_id=run_id,
                    error_code="sandbox_unavailable",
                    error_message=str(exc),
                )
            )
            raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc

        if resp.status_code >= 400:
            exc_type: type[SandboxError]
            if resp.status_code == 401:
                exc_type = SandboxUnauthorized
            elif resp.status_code == 400:
                exc_type = SandboxBadRequest
            elif resp.status_code == 404:
                exc_type = SandboxNotFound
            else:
                exc_type = SandboxExecFailed
            await self._fail(tool, trace_id, run_id, resp, exc_type, sink)

        result = response_cls.model_validate(resp.json())
        # completed 事件：read 带 returned 字节数 + 截断标记，write 带 bytes_written，
        # list/stat 无字节数语义，以 0 占位（exit_code=0 表示操作成功）。
        stdout_len = getattr(result, "returned", None) or getattr(result, "bytes_written", 0)
        truncated = bool(getattr(result, "truncated", False))
        await sink.emit(
            completed_event(
                tool=tool,
                trace_id=trace_id,
                run_id=run_id,
                exit_code=0,
                duration_ms=int((time.monotonic() - t0) * 1000),
                timeout=False,
                truncated=truncated,
                stdout_len=stdout_len,
                stderr_len=0,
            )
        )
        return result

    def _file_auth_headers(
        self,
        auth_ctx: AuthorizationContext,
        op: str,
        path: str,
    ) -> dict[str, str]:
        """签发文件操作凭证：file_op 绑定 {op, path}（exec 路径 body_hash 为 None，不绑定）。"""
        if not self._settings.sandbox_auth_enabled:
            return {}
        token = sign_authorization(
            self._settings.sandbox_auth_secret,
            auth_ctx,
            argv=None,
            stdin=None,
            ttl_seconds=self._settings.sandbox_auth_ttl_seconds,
            file_op=(op, path),
        )
        return {"Authorization": f"Bearer {token}"}

    def _auth_headers(
        self,
        auth_ctx: AuthorizationContext,
        argv: list[str] | None = None,
        stdin: str | None = None,
        *,
        code: str | None = None,
        language: str | None = None,
    ) -> dict[str, str]:
        """签发授权凭证。

        - exec 命令模式：传 argv/stdin，按 compute_body_hash 绑定。
        - exec 代码模式：传 code/language/stdin，按 compute_body_hash_code 绑定。
        - cancel：都不传，不绑定 body_hash。
        """
        if not self._settings.sandbox_auth_enabled:
            return {}
        token = sign_authorization(
            self._settings.sandbox_auth_secret,
            auth_ctx,
            argv=argv,
            stdin=stdin,
            code=code,
            language=language,
            ttl_seconds=self._settings.sandbox_auth_ttl_seconds,
        )
        return {"Authorization": f"Bearer {token}"}

    def _exec_auth_headers(
        self,
        auth_ctx: AuthorizationContext,
        request: ExecRequest,
    ) -> dict[str, str]:
        """按 code/argv 模式分流签发 exec 授权头。"""
        if request.code is not None:
            return self._auth_headers(
                auth_ctx,
                stdin=request.stdin,
                code=request.code,
                language=request.language,
            )
        return self._auth_headers(auth_ctx, request.argv, request.stdin)

    @staticmethod
    async def _read_error(resp: httpx.Response) -> str:
        """流式响应的错误体：先 aread 再复用 _extract_error。"""
        await resp.aread()
        _, message = SandboxClient._extract_error(resp)
        return message

    @staticmethod
    async def _iter_sse(resp: httpx.Response) -> AsyncIterator[tuple[str, str]]:
        """极简 SSE 解析：按空行分界，聚合 ``event:`` 与 ``data:`` 行（忽略 id/retry/注释/心跳）。"""
        event_type = "message"
        data_lines: list[str] = []
        async for line in resp.aiter_lines():
            if line == "":
                if data_lines:
                    yield event_type, "\n".join(data_lines)
                event_type = "message"
                data_lines = []
            elif line.startswith("event:"):
                event_type = line[len("event:") :].strip()
            elif line.startswith("data:"):
                data_lines.append(line[len("data:") :].lstrip())
            # 其余行（id:/retry:/注释/心跳）忽略
        if data_lines:
            yield event_type, "\n".join(data_lines)

    async def _fail(
        self,
        tool: str,
        trace_id: str | None,
        run_id: str | None,
        resp: httpx.Response,
        exc_type: type[SandboxError],
        sink: EventSink,
    ) -> NoReturn:
        code, message = self._extract_error(resp)
        await sink.emit(
            failed_event(
                tool=tool,
                trace_id=trace_id,
                run_id=run_id,
                error_code=code,
                error_message=message,
            )
        )
        raise exc_type(message)

    @staticmethod
    def _extract_error(resp: httpx.Response) -> tuple[str, str]:
        try:
            body = resp.json()
        except ValueError:
            return "exec_failed", resp.text
        err = body.get("error") if isinstance(body, dict) else None
        if isinstance(err, dict):
            return str(err.get("code", "exec_failed")), str(err.get("message", resp.text))
        return "exec_failed", resp.text