"""OpenHands Cloud API (V1) minimal client.

This file is intentionally:
- small (easy to copy into other repos)
- dependency-light (only `httpx`)
- opinionated in a helpful way (defaults to OpenHands Cloud)

Audience: AI agents.

The V1 API is hosted on the OpenHands app server under:
  {BASE_URL}/api/v1/...

Typical workflow for common operations:
1) Discover: GET /api/v1/users/me
2) List/search conversations: GET /api/v1/app-conversations/search
3) Start a conversation (creates sandbox): POST /api/v1/app-conversations
4) Monitor events for a conversation: GET /api/v1/conversation/{id}/events/search
5) (Optional) download trajectory: GET /api/v1/app-conversations/{id}/download

Note: Some operations happen against the *agent server* running inside a sandbox
(not the app server). Those endpoints use X-Session-API-Key instead of Bearer auth.

This client purposefully keeps responses as raw dicts/lists so agents can quickly
adapt it without strict schema maintenance.
"""

from __future__ import annotations

import argparse
import json
import os
import time
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import httpx


DEFAULT_BASE_URL = "https://app.all-hands.dev"
PREFERRED_API_KEY_ENV_VARS = ("OPENHANDS_CLOUD_API_KEY", "OPENHANDS_API_KEY")


# Start-task statuses observed in the wild. These may evolve, so keep this centralized.
START_TASK_TERMINAL_STATUSES = frozenset(
    {"READY", "ERROR", "FAILED", "CANCELLED", "DONE", "COMPLETED"}
)


# Safety cap for paging calls. Keeps responses small and consistent across clients.
AGENT_EVENTS_SEARCH_MAX_LIMIT = 100


@dataclass(frozen=True)
class OpenHandsAPIConfig:
    api_key: str
    base_url: str = DEFAULT_BASE_URL

    @property
    def api_v1_url(self) -> str:
        return f"{self.base_url.rstrip('/')}/api/v1"


class OpenHandsAPI:
    """Minimal OpenHands Cloud API client for the supported V1 API."""

    def __init__(self, api_key: str | None = None, base_url: str = DEFAULT_BASE_URL):
        resolved_key = api_key
        if not resolved_key:
            for env_name in PREFERRED_API_KEY_ENV_VARS:
                resolved_key = os.getenv(env_name)
                if resolved_key:
                    break
        if not resolved_key:
            env_list = ", ".join(PREFERRED_API_KEY_ENV_VARS)
            raise ValueError(f"Missing API key. Set one of: {env_list}, or pass api_key=...")

        self._cfg = OpenHandsAPIConfig(api_key=resolved_key, base_url=base_url.rstrip("/"))
        self._client = httpx.Client(
            headers={
                "Authorization": f"Bearer {self._cfg.api_key}",
                "Content-Type": "application/json",
            },
            timeout=30,
        )

    @property
    def base_url(self) -> str:
        return self._cfg.base_url

    @property
    def api_v1_url(self) -> str:
        return self._cfg.api_v1_url

    def close(self) -> None:
        self._client.close()

    # -----------------------------
    # App server endpoints (Bearer auth)
    # -----------------------------

    def users_me(self) -> dict[str, Any]:
        r = self._client.get(f"{self.api_v1_url}/users/me")
        r.raise_for_status()
        return r.json()

    def app_conversations_search(self, *, limit: int = 20) -> dict[str, Any]:
        limit = max(1, int(limit))
        r = self._client.get(
            f"{self.api_v1_url}/app-conversations/search", params={"limit": limit}
        )
        r.raise_for_status()
        return r.json()

    def app_conversations_count(self) -> dict[str, Any]:
        r = self._client.get(f"{self.api_v1_url}/app-conversations/count")
        r.raise_for_status()
        return r.json()

    def app_conversations_get_batch(self, *, ids: list[str]) -> list[dict[str, Any]]:
        if not ids:
            return []
        r = self._client.get(f"{self.api_v1_url}/app-conversations", params={"ids": ids})
        r.raise_for_status()
        return r.json()

    def app_conversation_get(self, conversation_id: str) -> dict[str, Any] | None:
        items = self.app_conversations_get_batch(ids=[conversation_id])
        return items[0] if items else None

    def sandboxes_search(self, *, limit: int = 20) -> dict[str, Any]:
        limit = max(1, int(limit))
        r = self._client.get(f"{self.api_v1_url}/sandboxes/search", params={"limit": limit})
        r.raise_for_status()
        return r.json()

    def sandbox_specs_search(self, *, limit: int = 20) -> dict[str, Any]:
        limit = max(1, int(limit))
        r = self._client.get(
            f"{self.api_v1_url}/sandbox-specs/search", params={"limit": limit}
        )
        r.raise_for_status()
        return r.json()

    def conversation_events_search(
        self, conversation_id: str, *, limit: int = 50
    ) -> dict[str, Any]:
        limit = max(1, int(limit))
        r = self._client.get(
            f"{self.api_v1_url}/conversation/{conversation_id}/events/search",
            params={"limit": limit},
        )
        r.raise_for_status()
        return r.json()

    def conversation_events_count(self, conversation_id: str) -> dict[str, Any]:
        r = self._client.get(f"{self.api_v1_url}/conversation/{conversation_id}/events/count")
        r.raise_for_status()
        return r.json()

    def app_conversation_start(
        self,
        *,
        initial_message: str,
        selected_repository: str | None = None,
        selected_branch: str | None = None,
        title: str | None = None,
        run: bool = True,
    ) -> dict[str, Any]:
        """Start a new V1 app conversation.

        WARNING: This typically creates a sandbox and may incur costs.

        In many deployments this endpoint is **asynchronous** and returns a **start-task** dict.
        Common fields:
        - `id`: the *start_task_id*
        - `app_conversation_id`: the id to use for `/download` and `/conversation/.../events/...`

        If `app_conversation_id` is missing from the initial response, fetch it via:
        - `GET /api/v1/app-conversations/start-tasks?ids=<start_task_id>`
          (see `app_conversation_start_task_get()` / `poll_start_task_until_ready()`).

        The payload structure here mirrors what the V1 app server expects:
        - initial_message.content is a list of content parts
        """

        payload: dict[str, Any] = {
            "initial_message": {
                "role": "user",
                # V1 expects `content` as a list of parts, even for a single text message.
                "content": [{"type": "text", "text": initial_message}],
                "run": bool(run),
            }
        }
        if selected_repository:
            payload["selected_repository"] = selected_repository
        if selected_branch:
            payload["selected_branch"] = selected_branch
        if title:
            payload["title"] = title

        r = self._client.post(f"{self.api_v1_url}/app-conversations", json=payload, timeout=120)
        r.raise_for_status()
        return r.json()

    def app_conversations_start_tasks_get_batch(self, *, ids: list[str]) -> list[dict[str, Any]]:
        if not ids:
            return []
        r = self._client.get(
            f"{self.api_v1_url}/app-conversations/start-tasks", params={"ids": ids}
        )
        r.raise_for_status()
        return r.json()

    def app_conversation_start_task_get(self, task_id: str) -> dict[str, Any] | None:
        items = self.app_conversations_start_tasks_get_batch(ids=[task_id])
        return items[0] if items else None

    def sandboxes_pause(self, sandbox_id: str) -> dict[str, Any]:
        r = self._client.post(f"{self.api_v1_url}/sandboxes/{sandbox_id}/pause", timeout=60)
        r.raise_for_status()
        return r.json()

    def sandboxes_resume(self, sandbox_id: str) -> dict[str, Any]:
        r = self._client.post(f"{self.api_v1_url}/sandboxes/{sandbox_id}/resume", timeout=60)
        r.raise_for_status()
        return r.json()

    def app_conversation_download_zip(
        self, app_conversation_id: str, *, output_file: str | Path
    ) -> dict[str, Any]:
        """Download a conversation trajectory zip to disk.

        Note: this endpoint expects the **app_conversation_id** (not the start-task id).
        """
        url = f"{self.api_v1_url}/app-conversations/{app_conversation_id}/download"
        r = self._client.get(url, timeout=60)
        r.raise_for_status()
        out = Path(output_file)
        out.write_bytes(r.content)
        return {
            "file": str(out),
            "size": len(r.content),
            "content_type": r.headers.get("content-type"),
        }

    def count_events_via_trajectory_zip(
        self,
        app_conversation_id: str,
        *,
        zip_file: str | Path,
        extract_dir: str | Path,
    ) -> dict[str, Any]:
        """Fallback event counting: download trajectory zip, extract, count event files.

        This is heavier than calling a count endpoint, but it is still a single API call and
        also gives you the full exported event payloads.

        Cleanup (optional): this helper writes a zip file and extracts JSON events. If you
        want to clean up afterwards, you can remove them, e.g.:

        - `zip_path.unlink(missing_ok=True)`
        - `shutil.rmtree(extract_path, ignore_errors=True)`

        Returns a small summary dict including `event_count`.
        """

        zip_path = Path(zip_file)
        extract_path = Path(extract_dir)

        download_meta = self.app_conversation_download_zip(app_conversation_id, output_file=zip_path)
        extract_path.mkdir(parents=True, exist_ok=True)

        with zipfile.ZipFile(zip_path, "r") as zf:
            zf.extractall(extract_path)

        event_count = len(list(extract_path.glob("event_*.json")))
        has_meta = (extract_path / "meta.json").exists()

        return {
            "event_count": event_count,
            "has_meta": has_meta,
            "zip": download_meta,
            "extract_dir": str(extract_path),
        }

    # -----------------------------
    # Agent server endpoints (X-Session-API-Key)
    # -----------------------------

    @staticmethod
    def agent_headers(session_api_key: str) -> dict[str, str]:
        return {"X-Session-API-Key": session_api_key, "Content-Type": "application/json"}


    @staticmethod
    def _agent_event_filter_params(
        *,
        timestamp_gte: str | None = None,
        timestamp_lt: str | None = None,
        kind: str | None = None,
        source: str | None = None,
        body: str | None = None,
    ) -> dict[str, Any]:
        params: dict[str, Any] = {}
        if timestamp_gte is not None:
            params["timestamp__gte"] = timestamp_gte
        if timestamp_lt is not None:
            params["timestamp__lt"] = timestamp_lt
        if kind is not None:
            params["kind"] = kind
        if source is not None:
            params["source"] = source
        if body is not None:
            params["body"] = body
        return params

    def agent_events_search(
        self,
        *,
        agent_server_url: str,
        session_api_key: str,
        conversation_id: str,
        limit: int = 50,
        sort_order: str | None = None,
        timestamp_gte: str | None = None,
        timestamp_lt: str | None = None,
        kind: str | None = None,
        source: str | None = None,
        body: str | None = None,
    ) -> dict[str, Any]:
        """Search events via the sandbox agent-server.

        Notes:
        - `limit` is capped at AGENT_EVENTS_SEARCH_MAX_LIMIT to avoid huge responses.
        - `sort_order` must be one of: "TIMESTAMP", "TIMESTAMP_DESC".
        - timestamp filters are passed as ISO-8601 strings (e.g. "2026-02-14T21:54:00Z").
          The server accepts both timezone-aware and naive datetimes.
        """

        url = f"{agent_server_url.rstrip('/')}/api/conversations/{conversation_id}/events/search"
        capped_limit = min(AGENT_EVENTS_SEARCH_MAX_LIMIT, max(1, int(limit)))
        params: dict[str, Any] = {"limit": capped_limit}
        if sort_order is not None:
            params["sort_order"] = sort_order
        params.update(
            self._agent_event_filter_params(
                timestamp_gte=timestamp_gte,
                timestamp_lt=timestamp_lt,
                kind=kind,
                source=source,
                body=body,
            )
        )

        r = httpx.get(
            url,
            headers=self.agent_headers(session_api_key),
            params=params,
            timeout=30,
        )
        r.raise_for_status()
        return r.json()

    def agent_events_count(
        self,
        *,
        agent_server_url: str,
        session_api_key: str,
        conversation_id: str,
        timestamp_gte: str | None = None,
        timestamp_lt: str | None = None,
        kind: str | None = None,
        source: str | None = None,
        body: str | None = None,
    ) -> int:
        """Count events via the sandbox agent-server.

        Timestamp filters are passed as ISO-8601 strings (e.g. "2026-02-14T21:54:00Z").
        """

        url = f"{agent_server_url.rstrip('/')}/api/conversations/{conversation_id}/events/count"
        params = self._agent_event_filter_params(
            timestamp_gte=timestamp_gte,
            timestamp_lt=timestamp_lt,
            kind=kind,
            source=source,
            body=body,
        )

        r = httpx.get(
            url,
            headers=self.agent_headers(session_api_key),
            params=params,
            timeout=30,
        )
        r.raise_for_status()
        return int(r.json())

    def agent_execute_bash(
        self,
        *,
        agent_server_url: str,
        session_api_key: str,
        command: str,
        cwd: str | None = None,
        timeout_s: int = 30,
    ) -> dict[str, Any]:
        url = f"{agent_server_url.rstrip('/')}/api/bash/execute_bash_command"
        payload: dict[str, Any] = {"command": command, "timeout": int(timeout_s)}
        if cwd:
            payload["cwd"] = cwd
        r = httpx.post(url, headers=self.agent_headers(session_api_key), json=payload, timeout=60)
        r.raise_for_status()
        return r.json()

    def agent_download_file(
        self,
        *,
        agent_server_url: str,
        session_api_key: str,
        path: str,
        output_file: str | Path,
    ) -> dict[str, Any]:
        p = path if path.startswith("/") else f"/{path}"
        url = f"{agent_server_url.rstrip('/')}/api/file/download{p}"
        r = httpx.get(url, headers=self.agent_headers(session_api_key), timeout=30)
        r.raise_for_status()
        out = Path(output_file)
        out.write_bytes(r.content)
        return {"file": str(out), "size": len(r.content)}

    def agent_upload_text_file(
        self,
        *,
        agent_server_url: str,
        session_api_key: str,
        path: str,
        content: str,
        content_type: str = "text/plain",
    ) -> dict[str, Any]:
        p = path if path.startswith("/") else f"/{path}"
        url = f"{agent_server_url.rstrip('/')}/api/file/upload{p}"
        filename = os.path.basename(p)
        headers = {"X-Session-API-Key": session_api_key}
        files = {"file": (filename, content.encode("utf-8"), content_type)}
        r = httpx.post(url, headers=headers, files=files, timeout=30)
        r.raise_for_status()
        return r.json() if r.text else {"success": True}

    # -----------------------------
    # Convenience helpers
    # -----------------------------

    def app_conversation_start_from_prompt_files(
        self,
        prompt_file: str | Path,
        *,
        selected_repository: str | None = None,
        selected_branch: str | None = None,
        title: str | None = None,
        append_file: str | Path | None = None,
        run: bool = True,
    ) -> dict[str, Any]:
        main_text = Path(prompt_file).read_text(encoding="utf-8")
        if append_file and Path(append_file).exists():
            tail = Path(append_file).read_text(encoding="utf-8")
            initial = f"{main_text}\n\n{tail}"
        else:
            initial = main_text

        return self.app_conversation_start(
            initial_message=initial,
            selected_repository=selected_repository,
            selected_branch=selected_branch,
            title=title,
            run=run,
        )


    @staticmethod
    def _start_task_status(task: dict[str, Any] | None) -> str:
        return str((task or {}).get("status") or "").upper()

    def poll_start_task_until_ready(
        self,
        task_id: str,
        *,
        timeout_s: int = 10 * 60,
        poll_interval_s: float = 2.0,
        backoff_factor: float = 1.5,
        max_interval_s: float = 10.0,
        max_polls: int | None = None,
    ) -> dict[str, Any]:
        """Poll a start-task until it reaches a terminal state.

        This is the async companion to `POST /api/v1/app-conversations`.

        It is intentionally *polite*:
        - sleeps between requests
        - uses exponential backoff (capped by `max_interval_s`)
        - supports `max_polls` to cap the total number of API calls

        Terminal statuses are defined in START_TASK_TERMINAL_STATUSES.

        Raises:
            TimeoutError: if the task doesn't reach a terminal state in time.
        """

        deadline = time.monotonic() + float(timeout_s)
        interval = max(0.25, float(poll_interval_s))
        factor = max(1.0, float(backoff_factor))
        max_interval = max(interval, float(max_interval_s))

        polls = 0
        last: dict[str, Any] | None = None

        while True:
            if max_polls is not None and polls >= int(max_polls):
                raise TimeoutError(
                    f"Start task {task_id} did not reach terminal state (max_polls={max_polls}, last={last})"
                )

            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TimeoutError(
                    f"Start task {task_id} did not reach terminal state in {timeout_s}s (last={last})"
                )

            last = self.app_conversation_start_task_get(task_id)
            polls += 1

            status = self._start_task_status(last)
            if status in START_TASK_TERMINAL_STATUSES:
                return last or {}

            sleep_s = min(interval, remaining)
            if sleep_s > 0:
                time.sleep(sleep_s)
            interval = min(max_interval, interval * factor)


OpenHandsV1API = OpenHandsAPI

def _cmd_search_conversations(args: argparse.Namespace) -> int:
    api = OpenHandsAPI(api_key=args.api_key, base_url=args.base_url)
    try:
        print(json.dumps(api.app_conversations_search(limit=args.limit), indent=2))
        return 0
    finally:
        api.close()


def _cmd_start_conversation(args: argparse.Namespace) -> int:
    api = OpenHandsAPI(api_key=args.api_key, base_url=args.base_url)
    try:
        resp = api.app_conversation_start_from_prompt_files(
            args.prompt_file,
            selected_repository=args.repo,
            selected_branch=args.branch,
            title=args.title,
            append_file=args.append_file,
            run=not args.no_run,
        )
        print(json.dumps(resp, indent=2))
        return 0
    finally:
        api.close()


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(prog="openhands_api.py")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p_search = sub.add_parser("search-conversations", help="GET /api/v1/app-conversations/search")
    p_search.add_argument(
        "--api-key",
        default=None,
        help="Defaults to OPENHANDS_CLOUD_API_KEY, then OPENHANDS_API_KEY",
    )
    p_search.add_argument("--base-url", default=DEFAULT_BASE_URL)
    p_search.add_argument("--limit", type=int, default=5)
    p_search.set_defaults(func=_cmd_search_conversations)

    p_start = sub.add_parser("start-conversation", help="POST /api/v1/app-conversations from a prompt file")
    p_start.add_argument(
        "--api-key",
        default=None,
        help="Defaults to OPENHANDS_CLOUD_API_KEY, then OPENHANDS_API_KEY",
    )
    p_start.add_argument("--base-url", default=DEFAULT_BASE_URL)
    p_start.add_argument("--prompt-file", required=True)
    p_start.add_argument("--append-file", default=None)
    p_start.add_argument("--repo", default=None)
    p_start.add_argument("--branch", default=None)
    p_start.add_argument("--title", default=None)
    p_start.add_argument("--no-run", action="store_true", help="If set, do not auto-run after sending initial message")
    p_start.set_defaults(func=_cmd_start_conversation)

    args = parser.parse_args(argv)
    return int(args.func(args))


if __name__ == "__main__":
    raise SystemExit(main())