"""OpenCode Environment Client.

Thin MCP client over the deployed ``opencode-openenv`` server. The server
exposes a single tool ``run_rollout`` that takes a task + LLM endpoint,
runs one OpenCode agent rollout in a fresh E2B sandbox, and returns a
JSON-serialized :class:`RolloutResult`.

Example::

    from opencode_env_server import OpenCodeEnv

    with OpenCodeEnv(base_url="https://adithya-s-k-opencode-openenv.hf.space") as env:
        env.reset()
        result = env.run_rollout(
            vllm_url="https://your-llm-host/v1",
            model="Qwen/Qwen3.5-4B",
            instruction="Write fizzbuzz.py in the current directory.",
            test_script="#!/bin/bash\\n...",
            task_id="fizzbuzz_001",
            mode="transparent_proxy",
            disable_thinking=True,
        )
        print(result.reward, len(result.proxy_turns))

Docker convenience::

    env = OpenCodeEnv.from_docker_image("opencode-openenv:latest")
    env.reset()
    ...
"""

from __future__ import annotations

import json
from typing import Any

from openenv.core.mcp_client import MCPToolClient

try:
    from .models import RolloutResult
except ImportError:
    from models import RolloutResult  # type: ignore


class OpenCodeEnv(MCPToolClient):
    """Client for the OpenCode OpenEnv server.

    Inherits MCP plumbing (``reset``, ``call_tool``, ``list_tools``,
    ``from_docker_image``, context-manager support) from
    :class:`MCPToolClient`. Adds :meth:`run_rollout` as a typed helper that
    deserializes the tool result into a :class:`RolloutResult`.
    """

    def run_rollout(
        self,
        *,
        vllm_url: str,
        model: str,
        instruction: str,
        test_script: str,
        task_id: str = "",
        setup_shell: str = "",
        upload_files: dict[str, str] | None = None,
        provider: str = "openai_compatible",
        api_key: str = "intercepted",
        mode: str = "transparent_proxy",
        disable_thinking: bool = False,
        max_tokens_cap: int = 4096,
        agent_timeout_s: float = 600.0,
    ) -> RolloutResult:
        """Typed helper around ``call_tool("run_rollout", ...)``."""

        raw = self.call_tool(
            "run_rollout",
            vllm_url=vllm_url,
            model=model,
            instruction=instruction,
            test_script=test_script,
            task_id=task_id,
            setup_shell=setup_shell,
            upload_files=upload_files or {},
            provider=provider,
            api_key=api_key,
            mode=mode,
            disable_thinking=disable_thinking,
            max_tokens_cap=max_tokens_cap,
            agent_timeout_s=agent_timeout_s,
        )
        payload = _extract_text(raw)
        return RolloutResult.model_validate_json(payload)


def _extract_text(result: Any) -> str:
    """Pull the text payload out of an MCP tool result shape.

    Handles three shapes MCPToolClient / call_tool may return:
    - the raw tool text (str)
    - a CallToolObservation-like object with ``.result.content[0].text``
    - a dict with ``content`` list containing ``{"text": ...}`` entries
    """

    if isinstance(result, str):
        return result

    # Object with attribute chain: obs.result.content[0].text
    inner = getattr(result, "result", None)
    if inner is not None:
        content = getattr(inner, "content", None)
        if content:
            first = content[0]
            text = getattr(first, "text", None)
            if isinstance(text, str):
                return text
            if isinstance(first, dict) and "text" in first:
                return first["text"]

    if isinstance(result, dict):
        content = result.get("content")
        if isinstance(content, list) and content:
            first = content[0]
            if isinstance(first, dict) and "text" in first:
                return first["text"]
        nested = result.get("result")
        if isinstance(nested, dict):
            content = nested.get("content")
            if isinstance(content, list) and content:
                first = content[0]
                if isinstance(first, dict) and "text" in first:
                    return first["text"]
        return json.dumps(result, default=str)

    # Object with .content directly
    content = getattr(result, "content", None)
    if content:
        first = content[0]
        text = getattr(first, "text", None)
        if isinstance(text, str):
            return text

    return str(result)