import json
import os
import shlex
from pathlib import Path
from typing import Any

from harbor.agents.installed.base import BaseInstalledAgent, ExecInput
from harbor.agents.utils import get_api_key_var_names_from_model_name
from harbor.models.agent.context import AgentContext
from harbor.models.agent.name import AgentName
from harbor.models.agent.trajectory_config import TrajectoryConfig
from harbor.models.trajectories import (
    Agent,
    FinalMetrics,
    Metrics,
    Observation,
    ObservationResult,
    Step,
    ToolCall,
    Trajectory,
)
from harbor.models.trial.paths import EnvironmentPaths


class OpenHands(BaseInstalledAgent):
    """
    The OpenHands agent uses the All Hands AI OpenHands tool to solve tasks.
    """

    SUPPORTS_ATIF: bool = True

    def __init__(
        self,
        disable_tool_calls: bool = False,
        trajectory_config: TrajectoryConfig | None = None,
        reasoning_effort: str | None = "high",
        *args,
        **kwargs,
    ):
        """
        Initialize OpenHands agent.

        Args:
            disable_tool_calls: Whether to disable tool calls.
                If true, native function calling will be disabled.
            trajectory_config: Optional TrajectoryConfig containing trajectory-related configurations.
                Available options:
                - raw_content (bool): If True, read trajectory from raw completions folder instead
                  of parsed event logs. Useful for preserving raw LLM responses. (default: False)
            reasoning_effort: Reasoning effort level for the model (low, medium, high).
        """
        # Extract api_base and model_info before passing kwargs to super()
        self._api_base = kwargs.pop("api_base", None)
        self._model_info = kwargs.pop("model_info", None)
        super().__init__(*args, **kwargs)
        self._reasoning_effort = reasoning_effort
        self._disable_tool_calls = disable_tool_calls
        self._git_version = kwargs.get("git_version", None)
        self._trajectory_config = trajectory_config or {}
        self._save_raw_content_in_trajectory = self._trajectory_config.get(
            "raw_content", False
        )

        # Validate that raw_content requires disable_tool_calls
        if self._save_raw_content_in_trajectory and not self._disable_tool_calls:
            raise ValueError(
                "When trajectory_config['raw_content'] is True, disable_tool_calls must also be True. "
                "Raw content mode is designed to preserve text-based tool invocation format (e.g., <execute_bash>) "
                "rather than native function calling format. Native function calling does not appear in LLM response's "
                "message content field, rather, it is structured in the tool_calls field, thus there is no easy and "
                "standard way to persist it in the trajectory in its raw/original format. Note that even OpenHands may not "
                "use native function calling for some models, disable_tool_calls must still be set explicitly in raw_content mode."
            )

    @staticmethod
    def name() -> str:
        return AgentName.OPENHANDS.value

    @property
    def _trajectory_path(self) -> Path:
        return EnvironmentPaths.agent_dir / "trajectory.json"

    def _get_session_dir(self) -> Path | None:
        """Get the single session directory."""
        sessions_dir = self.logs_dir / "sessions"
        if not sessions_dir.exists():
            return None

        session_dirs = [d for d in sessions_dir.iterdir() if d.is_dir()]
        if not session_dirs:
            return None

        # Sanity check: there should be exactly one session
        assert len(session_dirs) == 1, (
            f"Expected exactly 1 session, found {len(session_dirs)}"
        )
        return session_dirs[0]

    def _extract_version_and_metadata(
        self, events_dir: Path, include_tool_definitions: bool = True
    ) -> tuple[str, dict[str, Any] | None, list[dict[str, Any]] | None]:
        """Extract version, extra metadata, and optionally tool definitions from event files.

        Args:
            events_dir: Directory containing event files
            include_tool_definitions: Whether to extract tool definitions (default: True)

        Returns:
            Tuple of (version, extra, tool_definitions)
        """
        version = "unknown"
        extra: dict[str, Any] | None = None
        tool_definitions: list[dict[str, Any]] | None = None

        event_files = sorted(events_dir.glob("*.json"), key=lambda p: int(p.stem))
        for event_file in event_files:
            try:
                with open(event_file, "r") as f:
                    event = json.load(f)
                    if "args" in event:
                        if "openhands_version" in event["args"]:
                            version = event["args"]["openhands_version"]

                        # Extract extra info
                        extra_data = {
                            key: event["args"][key]
                            for key in ["agent_class"]
                            if key in event["args"]
                        }
                        if extra_data:
                            extra = extra_data

                        # Extract tool definitions if requested and available
                        if include_tool_definitions and "tools" in event["args"]:
                            tool_definitions = event["args"]["tools"]

                        if version != "unknown":
                            break
            except Exception as e:
                print(f"Warning: Could not read event file {event_file}: {e}")

        return version, extra, tool_definitions

    def _convert_event_to_step(self, event: dict[str, Any], step_id: int) -> Step:
        """Convert an OpenHands event to ATIF trajectory step format using Pydantic models."""
        # Map OpenHands source to ATIF-compliant source
        source = event.get("source", "")
        message = event.get("message", "")

        # Determine if this is a system message based on various heuristics
        # Check if this is a system prompt (action="system" regardless of source)
        if event.get("action") == "system":
            source = "system"
        elif source == "environment":
            # Map environment to system as per ATIF spec
            source = "system"
        elif source == "user" and message:
            # Heuristic: System-generated messages often start with specific patterns
            # These are internal OpenHands status messages, not actual user input
            system_message_patterns = [
                "Retrieving content for:",
                "Added workspace context",
                "Loading workspace",
                "Initializing",
            ]
            if any(message.startswith(pattern) for pattern in system_message_patterns):
                source = "system"

        # Extract timestamp
        timestamp = event.get("timestamp", "")

        # Handle tool calls for agent actions (only for agent source)
        tool_calls: list[ToolCall] | None = None
        if event.get("source") == "agent" and "tool_call_metadata" in event:
            metadata = event["tool_call_metadata"]
            arguments: dict[str, Any] = {}

            # Extract arguments if available from model_response
            if "model_response" in metadata:
                model_resp = metadata["model_response"]
                if "choices" in model_resp and len(model_resp["choices"]) > 0:
                    choice = model_resp["choices"][0]
                    if "message" in choice and "tool_calls" in choice["message"]:
                        tc_list = choice["message"]["tool_calls"]
                        if tc_list and len(tc_list) > 0:
                            if (
                                "function" in tc_list[0]
                                and "arguments" in tc_list[0]["function"]
                            ):
                                try:
                                    # Parse arguments if they're a JSON string
                                    args = tc_list[0]["function"]["arguments"]
                                    if isinstance(args, str):
                                        arguments = json.loads(args)
                                    else:
                                        arguments = args
                                except (json.JSONDecodeError, KeyError):
                                    arguments = {}

            tool_calls = [
                ToolCall(
                    tool_call_id=metadata.get("tool_call_id", ""),
                    function_name=metadata.get("function_name", ""),
                    arguments=arguments,
                )
            ]

        # Handle observations (tool call results) - only for agent source
        observation: Observation | None = None
        if (
            "observation" in event
            and event.get("cause") is not None
            and event.get("source") == "agent"
        ):
            # Link observation to tool call if available
            source_call_id = None
            if "tool_call_metadata" in event:
                source_call_id = event["tool_call_metadata"].get("tool_call_id")

            # Note: We're not storing OpenHands-specific metadata (type, extras, success)
            # in the observation since the current Pydantic model doesn't have an extra field.
            # This could be added in the future if needed.
            observation = Observation(
                results=[
                    ObservationResult(
                        source_call_id=source_call_id,
                        content=event.get("content", ""),
                    )
                ]
            )

        # Add metrics if available (only for agent source)
        metrics: Metrics | None = None
        if "llm_metrics" in event and event.get("source") == "agent":
            llm_metrics = event["llm_metrics"]
            accumulated_usage = llm_metrics.get("accumulated_token_usage", {})
            prompt_tokens = accumulated_usage.get("prompt_tokens", 0)
            completion_tokens = accumulated_usage.get("completion_tokens", 0)
            cached_tokens = accumulated_usage.get("cache_read_tokens", 0)
            cost_usd = llm_metrics.get("accumulated_cost", 0)

            metrics = Metrics(
                prompt_tokens=prompt_tokens if prompt_tokens > 0 else None,
                completion_tokens=completion_tokens if completion_tokens > 0 else None,
                cached_tokens=cached_tokens if cached_tokens > 0 else None,
                cost_usd=cost_usd if cost_usd > 0 else None,
            )

        return Step(
            step_id=step_id,
            timestamp=timestamp,
            source=source,
            message=message,
            tool_calls=tool_calls,
            observation=observation,
            metrics=metrics,
        )

    def _convert_events_to_trajectory(self, events_dir: Path) -> Trajectory | None:
        """
        Convert OpenHands event files to ATIF trajectory format.

        This method reads from the events folder which contains structured tool calls
        no matter if function calling is enabled or disabled for that model.

        Args:
            events_dir: Directory containing event files

        Returns:
            Trajectory object with parsed tool calls

        Raises:
            ValueError: If events directory doesn't exist or contains no event files
        """
        # Read all event files
        event_files = sorted(events_dir.glob("*.json"), key=lambda p: int(p.stem))
        events = []
        for event_file in event_files:
            try:
                with open(event_file, "r") as f:
                    events.append(json.load(f))
            except Exception as e:
                print(f"Warning: Could not read event file {event_file}: {e}")

        if not events:
            return None

        # Extract version and metadata from events
        version, extra, tool_definitions = self._extract_version_and_metadata(
            events_dir
        )

        # Generate session_id from events directory path
        # Use the parent directory name (session folder name) as session_id
        session_id = events_dir.parent.name

        # Convert events to steps (step_id starts from 1 per ATIF spec)
        # Note: OpenHands stores accumulated metrics in each event, but ATIF expects
        # per-step deltas. We'll need to calculate deltas after conversion.
        steps: list[Step] = []
        step_counter = 1
        for event in events:
            step = self._convert_event_to_step(event, step_counter)

            # OpenHands trajectories contain bookkeeping system steps
            # that don't contain meaningful information as trajectory steps.
            if step.message or step.tool_calls or step.observation:
                # Check if this step should be merged with the previous step
                # OpenHands sometimes emits two events for the same action:
                # 1. Agent action with tool_call but no observation
                # 2. Same agent action with observation (sometimes repeating the tool_call)
                # We merge these into a single step by checking if the tool_call_id matches
                should_merge = False
                if (
                    steps
                    and step.source == "agent"
                    and step.observation
                    and step.tool_calls
                    and steps[-1].source == "agent"
                    and steps[-1].tool_calls
                    and not steps[-1].observation
                ):
                    # Check if the tool_call_ids match
                    prev_call_id = steps[-1].tool_calls[0].tool_call_id
                    curr_call_id = step.tool_calls[0].tool_call_id
                    if prev_call_id == curr_call_id:
                        should_merge = True

                if should_merge:
                    # Merge this step into the previous step
                    # Add observation to previous step and preserve message if present
                    steps[-1].observation = step.observation
                    if step.message and not steps[-1].message:
                        steps[-1].message = step.message
                    # Don't increment step_counter since we're merging
                else:
                    steps.append(step)
                    step_counter += 1

        # Convert accumulated metrics to per-step deltas
        # OpenHands events contain accumulated_token_usage, but ATIF steps should
        # contain per-step token usage. Calculate deltas between consecutive steps.
        prev_prompt = 0
        prev_completion = 0
        prev_cached = 0
        prev_cost = 0.0

        for step in steps:
            if step.metrics:
                # Get current accumulated values
                curr_prompt = step.metrics.prompt_tokens or 0
                curr_completion = step.metrics.completion_tokens or 0
                curr_cached = step.metrics.cached_tokens or 0
                curr_cost = step.metrics.cost_usd or 0.0

                # Calculate delta (this step's usage only)
                delta_prompt = curr_prompt - prev_prompt
                delta_completion = curr_completion - prev_completion
                delta_cached = curr_cached - prev_cached
                delta_cost = curr_cost - prev_cost

                # Update step metrics with deltas
                step.metrics.prompt_tokens = delta_prompt if delta_prompt > 0 else None
                step.metrics.completion_tokens = (
                    delta_completion if delta_completion > 0 else None
                )
                step.metrics.cached_tokens = delta_cached if delta_cached > 0 else None
                step.metrics.cost_usd = delta_cost if delta_cost > 0 else None

                # Update previous values for next iteration
                prev_prompt = curr_prompt
                prev_completion = curr_completion
                prev_cached = curr_cached
                prev_cost = curr_cost

        # Extract final metrics from last step with metrics
        final_metrics: FinalMetrics | None = None
        for event in reversed(events):
            if "llm_metrics" in event:
                accumulated_usage = event["llm_metrics"].get(
                    "accumulated_token_usage", {}
                )
                total_prompt_tokens = accumulated_usage.get("prompt_tokens", 0)
                total_completion_tokens = accumulated_usage.get("completion_tokens", 0)
                total_cached_tokens = accumulated_usage.get("cache_read_tokens", 0)
                total_cost_usd = event["llm_metrics"].get("accumulated_cost", 0)

                final_metrics = FinalMetrics(
                    total_prompt_tokens=total_prompt_tokens
                    if total_prompt_tokens > 0
                    else None,
                    total_completion_tokens=total_completion_tokens
                    if total_completion_tokens > 0
                    else None,
                    total_cached_tokens=total_cached_tokens
                    if total_cached_tokens > 0
                    else None,
                    total_cost_usd=total_cost_usd if total_cost_usd > 0 else None,
                )
                break

        # Build trajectory using Pydantic models
        trajectory = Trajectory(
            schema_version="ATIF-v1.5",
            session_id=session_id,
            agent=Agent(
                name="openhands",
                version=version,
                tool_definitions=tool_definitions,
                extra=extra,
            ),
            steps=steps,
            final_metrics=final_metrics,
        )

        return trajectory

    def _convert_completions_to_trajectory(
        self, completions_dir: Path, events_dir: Path
    ) -> Trajectory | None:
        """Convert OpenHands completion files to ATIF trajectory format with raw LLM responses.

        This method reads from the completions folder which preserves the original format of LLM responses
        when function calling is disabled or not available for that model.

        Args:
            completions_dir: Directory containing completion JSON files
            events_dir: Directory containing event files (used to extract tool definitions when needed)

        Returns:
            Trajectory object, or None if completions are not available
        """
        if not completions_dir.exists():
            print(
                f"Warning: Completions directory does not exist: {completions_dir}. "
                "Cannot generate trajectory when raw_content=True."
            )
            return None

        # Read all completion files sorted by timestamp in filename
        # Filename format: provider__model-name-timestamp.json
        # Example: anthropic__claude-sonnet-4-5-1765823587.968077.json
        def get_timestamp(path: Path) -> float:
            """Extract timestamp from completion filename."""
            try:
                # Split by '-' and get the last part before .json
                parts = path.stem.split("-")
                # The timestamp is typically at the end: timestamp1.timestamp2
                timestamp_str = parts[-1]
                return float(timestamp_str)
            except (ValueError, IndexError):
                return 0.0

        completion_files = sorted(
            completions_dir.glob("*.json"),
            key=get_timestamp,
        )

        if not completion_files:
            print(
                f"Warning: No completion files found in {completions_dir}. "
                "Cannot generate trajectory when raw_content=True."
            )
            return None

        # Extract session info from first completion
        try:
            with open(completion_files[0], "r") as f:
                first_completion = json.load(f)
        except Exception as e:
            print(f"Warning: Could not read first completion file: {e}")
            return None

        # Determine if using native function calling by checking if tools are in kwargs
        using_native_function_calling = "tools" in first_completion.get("kwargs", {})

        # Extract version and metadata from events
        # Only include tool_definitions when using native function calling
        # When using text-based tool invocation (disable_tool_calls=True, or
        # this model does not support function calling, or OpenHands wrongly assumes
        # it does not support function calling), the tools are already defined inline
        # in the system prompt
        version, extra, tool_definitions = self._extract_version_and_metadata(
            events_dir, include_tool_definitions=using_native_function_calling
        )

        # Generate session_id from completions directory path
        session_id = completions_dir.parent.parent.name

        # Build steps from completions
        steps: list[Step] = []
        step_counter = 1
        total_prompt_tokens = 0
        total_completion_tokens = 0
        total_cached_tokens = 0
        total_cost_usd = 0.0

        for completion_idx, completion_file in enumerate(completion_files):
            try:
                with open(completion_file, "r") as f:
                    completion = json.load(f)
            except Exception as e:
                print(f"Warning: Could not read completion file {completion_file}: {e}")
                continue

            messages = completion.get("messages", [])
            fncall_messages = completion.get("fncall_messages", [])
            response = completion.get("response", {})
            timestamp_float = completion.get("timestamp", 0)

            # Convert timestamp to ISO format
            from datetime import datetime, timezone

            timestamp = datetime.fromtimestamp(
                timestamp_float, tz=timezone.utc
            ).isoformat()

            # For the FIRST completion: extract system and user messages from messages array
            # For SUBSEQUENT completions: extract observation from fncall_messages
            if completion_idx == 0:
                # First completion: process system and user messages
                for msg in messages:
                    role = msg.get("role", "")
                    if role == "system":
                        # Add system message as first step
                        content_parts = msg.get("content", [])
                        if isinstance(content_parts, list) and len(content_parts) > 0:
                            system_content = content_parts[0].get("text", "")
                            steps.append(
                                Step(
                                    step_id=step_counter,
                                    timestamp=timestamp,
                                    source="system",
                                    message=system_content,
                                )
                            )
                            step_counter += 1
                    elif role == "user":
                        # Add user message
                        content_parts = msg.get("content", [])
                        if isinstance(content_parts, list) and len(content_parts) > 0:
                            user_content = content_parts[0].get("text", "")
                            steps.append(
                                Step(
                                    step_id=step_counter,
                                    timestamp=timestamp,
                                    source="user",
                                    message=user_content,
                                )
                            )
                            step_counter += 1
            else:
                # Subsequent completions: extract observation from fncall_messages
                # The observation is the last "tool" role message in fncall_messages
                for msg in reversed(fncall_messages):
                    if msg.get("role") == "tool":
                        # Extract observation content
                        content_data = msg.get("content", [])
                        if isinstance(content_data, list) and len(content_data) > 0:
                            observation_content = content_data[0].get("text", "")
                        elif isinstance(content_data, str):
                            observation_content = content_data
                        else:
                            observation_content = ""

                        # Attach observation to the previous agent step
                        # In text-based tool invocation mode, we don't have structured tool_calls,
                        # so we don't set source_call_id (leaving it as None is valid per ATIF spec)
                        if steps and steps[-1].source == "agent":
                            steps[-1].observation = Observation(
                                results=[
                                    ObservationResult(
                                        source_call_id=None,
                                        content=observation_content,
                                    )
                                ]
                            )
                        break

            # Process response - extract raw content and metrics
            choices = response.get("choices", [])
            if not choices:
                continue

            message = choices[0].get("message", {})
            content = message.get("content") or ""

            # Extract usage metrics
            usage = response.get("usage", {})
            prompt_tokens = usage.get("prompt_tokens", 0)
            completion_tokens = usage.get("completion_tokens", 0)
            # Handle different cache token field names
            cached_tokens = usage.get("cache_read_tokens", 0) or usage.get(
                "cache_read_input_tokens", 0
            )

            # Get cost from completion metadata
            cost = completion.get("cost", 0.0)

            # Update totals
            total_prompt_tokens += prompt_tokens
            total_completion_tokens += completion_tokens
            total_cached_tokens += cached_tokens
            total_cost_usd += cost

            # Create metrics object
            metrics = Metrics(
                prompt_tokens=prompt_tokens if prompt_tokens > 0 else None,
                completion_tokens=completion_tokens if completion_tokens > 0 else None,
                cached_tokens=cached_tokens if cached_tokens > 0 else None,
                cost_usd=cost if cost > 0 else None,
            )

            # Add agent response step with raw content (no tool_calls parsing)
            steps.append(
                Step(
                    step_id=step_counter,
                    timestamp=timestamp,
                    source="agent",
                    message=content,
                    metrics=metrics,
                )
            )
            step_counter += 1

        # Build final metrics
        final_metrics = FinalMetrics(
            total_prompt_tokens=total_prompt_tokens
            if total_prompt_tokens > 0
            else None,
            total_completion_tokens=total_completion_tokens
            if total_completion_tokens > 0
            else None,
            total_cached_tokens=total_cached_tokens
            if total_cached_tokens > 0
            else None,
            total_cost_usd=total_cost_usd if total_cost_usd > 0 else None,
        )

        # Build trajectory
        trajectory = Trajectory(
            schema_version="ATIF-v1.5",
            session_id=session_id,
            agent=Agent(
                name="openhands",
                version=version,
                tool_definitions=tool_definitions,
                extra=extra,
            ),
            steps=steps,
            final_metrics=final_metrics,
        )

        return trajectory

    def populate_context_post_run(self, context: AgentContext) -> None:
        """
        Populate context after agent run completes or times out.

        This method:
        1. Gets the session directory
        2. Converts individual event/completion files to a single trajectory.json file
        3. Populates context with token usage and cost information

        When raw_content=True in trajectory_config:
        - Reads from agent/completions folder which preserves raw LLM responses
        - Requires completions to be available (raises error if not)

        When raw_content=False (default):
        - Reads from events folder which contains parsed tool calls
        """
        # Get the session directory
        session_dir = self._get_session_dir()
        if not session_dir:
            print("No session directory found")
            return

        events_dir = session_dir / "events"
        if not events_dir.exists():
            print(f"Events directory {events_dir} does not exist")
            return

        # Convert to trajectory based on raw_content setting
        trajectory: Trajectory | None = None
        if self._save_raw_content_in_trajectory:
            # Use completions folder for raw LLM responses
            completions_dir = self.logs_dir / "completions"
            trajectory = self._convert_completions_to_trajectory(
                completions_dir, events_dir
            )
        else:
            # Use events folder for parsed tool calls
            trajectory = self._convert_events_to_trajectory(events_dir)

        if not trajectory:
            print("Failed to convert to trajectory")
            return

        trajectory_path = self.logs_dir / "trajectory.json"
        try:
            with open(trajectory_path, "w") as f:
                json.dump(trajectory.to_json_dict(), f, indent=2)
        except Exception as e:
            print(f"Failed to write trajectory file: {e}")
            return

        # Populate context from trajectory Pydantic model
        if trajectory.final_metrics:
            metrics = trajectory.final_metrics
            context.cost_usd = metrics.total_cost_usd
            context.n_input_tokens = metrics.total_prompt_tokens or 0
            context.n_cache_tokens = metrics.total_cached_tokens or 0
            context.n_output_tokens = metrics.total_completion_tokens or 0
        else:
            print("No final_metrics found in trajectory")

    @property
    def _template_variables(self) -> dict[str, str]:
        """
        Provide template variables for OpenHands setup script.

        Includes version and git_version for development builds.
        """
        variables = dict()

        if self._git_version:
            variables["git_version"] = self._git_version
        elif self._version:
            variables["version"] = self._version

        return variables

    @property
    def _install_agent_template_path(self) -> Path:
        return Path(__file__).parent / "install-openhands.sh.j2"

    def _build_mcp_config_toml(self) -> str | None:
        """Build OpenHands config.toml content with MCP server configuration.

        Returns ``None`` when no MCP servers are configured.
        """
        if not self.mcp_servers:
            return None

        sse_entries: list[str] = []
        shttp_entries: list[str] = []
        stdio_entries: list[str] = []

        for server in self.mcp_servers:
            if server.transport == "sse":
                sse_entries.append(f'{{url = "{server.url}"}}')
            elif server.transport == "streamable-http":
                shttp_entries.append(f'{{url = "{server.url}"}}')
            elif server.transport == "stdio":
                args_str = ", ".join(f'"{a}"' for a in server.args)
                stdio_entries.append(
                    f'{{name = "{server.name}", command = "{server.command}", args = [{args_str}]}}'
                )

        lines = ["[mcp]"]
        if sse_entries:
            lines.append(f"sse_servers = [{', '.join(sse_entries)}]")
        if shttp_entries:
            lines.append(f"shttp_servers = [{', '.join(shttp_entries)}]")
        if stdio_entries:
            lines.append(f"stdio_servers = [{', '.join(stdio_entries)}]")

        return "\n".join(lines) + "\n"

    def create_run_agent_commands(self, instruction: str) -> list[ExecInput]:
        escaped_instruction = shlex.quote(instruction)

        env = {}

        is_hosted_vllm_model = (
            self.model_name is not None
            and self.model_name.lower().startswith("hosted_vllm/")
        )

        # Handle LLM API key with fallback logic
        # For hosted_vllm models (local vLLM servers), use a dummy key if none set
        if "LLM_API_KEY" in os.environ:
            env["LLM_API_KEY"] = os.environ["LLM_API_KEY"]
        elif is_hosted_vllm_model:
            # Local vLLM servers don't validate API keys, use dummy
            env["LLM_API_KEY"] = "dummy-key-for-local-vllm"
        else:
            if self.model_name:
                try:
                    api_key_vars = get_api_key_var_names_from_model_name(
                        self.model_name
                    )
                    if len(api_key_vars) == 1:
                        api_key_var = api_key_vars[0]
                        if api_key_var in os.environ:
                            env["LLM_API_KEY"] = os.environ[api_key_var]
                        else:
                            raise ValueError(
                                f"Unset API variable found for model {self.model_name}. "
                                f"Please set {api_key_var} or LLM_API_KEY environment variable"
                            )
                    else:
                        for api_key_var in api_key_vars:
                            if api_key_var in os.environ:
                                # Use model-agnostic variables received by OpenHands
                                oh_api_key_var = api_key_var.replace("AZURE_", "LLM_")
                                env[oh_api_key_var] = os.environ[api_key_var]
                            else:
                                raise ValueError(
                                    f"Unset API variable found for model {self.model_name}. "
                                    f"Please set {api_key_var} or LLM_API_KEY environment variable"
                                )
                except ValueError as e:
                    raise ValueError(
                        f"Unable to determine API key for model {self.model_name}: {e}. "
                        "Please set LLM_API_KEY environment variable as fallback"
                    )
            else:
                raise ValueError(
                    "No LLM API key found and no model specified. "
                    "Please set LLM_API_KEY environment variable or specify a model"
                )

        # Set model name
        if self.model_name:
            env["LLM_MODEL"] = self.model_name
        elif "LLM_MODEL" in os.environ:
            env["LLM_MODEL"] = os.environ["LLM_MODEL"]
        elif "ANTHROPIC_MODEL" in os.environ:
            env["LLM_MODEL"] = os.environ["ANTHROPIC_MODEL"]
        else:
            raise ValueError(
                "No LLM model found. Please set LLM_MODEL environment variable "
                "or specify it in the CLI"
            )

        if "LLM_BASE_URL" in os.environ:
            env["LLM_BASE_URL"] = os.environ["LLM_BASE_URL"]
        elif self._api_base:
            env["LLM_BASE_URL"] = self._api_base
        else:
            for candidate in (
                "HOSTED_VLLM_API_BASE",
                "VLLM_API_BASE",
                "OPENAI_API_BASE",
            ):
                if candidate in os.environ:
                    env["LLM_BASE_URL"] = os.environ[candidate]
                    break

        if "LLM_API_VERSION" in os.environ:
            env["LLM_API_VERSION"] = os.environ["LLM_API_VERSION"]

        # Pass model_info token limits to OpenHands (critical for hosted_vllm models)
        # Without these, OpenHands uses LiteLLM defaults which may be too restrictive
        if self._model_info:
            if self._model_info.get("max_input_tokens"):
                env["LLM_MAX_INPUT_TOKENS"] = str(self._model_info["max_input_tokens"])
            if self._model_info.get("max_output_tokens"):
                env["LLM_MAX_OUTPUT_TOKENS"] = str(
                    self._model_info["max_output_tokens"]
                )

        # Set reasoning_effort from kwargs or environment variable
        reasoning_effort = self._reasoning_effort
        if reasoning_effort is not None:
            env["LLM_REASONING_EFFORT"] = str(reasoning_effort)

        # Disable browsing and prompt extensions
        env["AGENT_ENABLE_PROMPT_EXTENSIONS"] = "false"
        env["AGENT_ENABLE_BROWSING"] = "false"
        env["ENABLE_BROWSER"] = "false"

        # Other sandbox settings
        env["SANDBOX_ENABLE_AUTO_LINT"] = "true"
        env["SKIP_DEPENDENCY_CHECK"] = "1"
        env["RUN_AS_OPENHANDS"] = "false"
        env["RUNTIME"] = "local"

        # trajectory saving
        # note this trajectory is of openhands format, not ATIF. Also, it is only
        # saved at the end of the run, not during each step.
        env["SAVE_TRAJECTORY_PATH"] = "/logs/agent/openhands.trajectory.json"

        # logging
        env["FILE_STORE"] = "local"
        env["FILE_STORE_PATH"] = "/logs/agent/"
        env["LLM_LOG_COMPLETIONS"] = "true"
        env["LLM_LOG_COMPLETIONS_FOLDER"] = "/logs/agent/completions/"

        # disable native function calling
        if self._disable_tool_calls:
            env["LLM_NATIVE_TOOL_CALLING"] = "false"

        # Handle any other OpenHands environment variables
        for key, value in os.environ.items():
            if key.startswith("OPENHANDS_"):
                env[key.replace("OPENHANDS_", "")] = value

        # Write MCP config.toml if MCP servers are configured
        mcp_config = self._build_mcp_config_toml()
        config_file_path = "~/.openhands/config.toml"

        exec_inputs: list[ExecInput] = []
        if mcp_config:
            escaped_config = shlex.quote(mcp_config)
            exec_inputs.append(
                ExecInput(
                    command=f"mkdir -p ~/.openhands && echo {escaped_config} > {config_file_path}",
                    env=env,
                )
            )

        commands = [
            # Use current directory as workspace
            "SANDBOX_VOLUMES=${PWD}:/workspace:rw",
            "/opt/openhands-venv/bin/python -m openhands.core.main",
            f"--task={escaped_instruction}",
        ]

        if mcp_config:
            commands.append(f"--config-file={config_file_path}")

        exec_inputs.append(
            ExecInput(
                command=" ".join(commands)
                + " 2>&1 </dev/null | stdbuf -oL tee /logs/agent/openhands.txt",
                env=env,
            )
        )

        return exec_inputs
