import requests
import time
import json
import base64
import binascii
from urllib.parse import quote
from typing import Optional, Dict, Any, Generator, Callable, List
from .exceptions import (
    InstaVMError, AuthenticationError, SessionError, ExecutionError,
    NetworkError, RateLimitError, BrowserError, BrowserSessionError,
    BrowserInteractionError, BrowserTimeoutError, BrowserNavigationError,
    ElementNotFoundError, QuotaExceededError, UnsupportedOperationError
)


def _encode_path_segment(value: Any) -> str:
    """Encode one URL path segment to avoid path traversal and delimiter injection."""
    segment = str(value)
    if segment in (".", ".."):
        raise ValueError("Path traversal not allowed in path segment")
    return quote(segment, safe="")


def _encode_proxy_path(path: str) -> str:
    """Encode a slash-delimited proxy path while preserving segment boundaries."""
    segments: List[str] = []
    for segment in path.lstrip("/").split("/"):
        if not segment or segment == ".":
            continue
        if segment == "..":
            raise ValueError("Path traversal not allowed in proxy path")
        segments.append(quote(segment, safe=""))
    return "/".join(segments)


class BrowserSession:
    """Represents a browser session for automation"""

    def __init__(self, session_id: str, client: 'InstaVM'):
        self.session_id = session_id
        self.client = client
        self._active = True

    def navigate(self, url: str, wait_timeout: int = 30000) -> Dict[str, Any]:
        """Navigate to a URL"""
        return self.client.browser_navigate(url, self.session_id, wait_timeout)

    def click(self, selector: str, force: bool = False, timeout: int = 30000) -> Dict[str, Any]:
        """Click an element"""
        return self.client.browser_click(selector, self.session_id, force, timeout)

    def type(self, selector: str, text: str, delay: int = 100, timeout: int = 30000) -> Dict[str, Any]:
        """Type text into an element"""
        return self.client.browser_type(selector, text, self.session_id, delay, timeout)

    def fill(self, selector: str, value: str, timeout: int = 30000) -> Dict[str, Any]:
        """Fill a form field"""
        return self.client.browser_fill(selector, value, self.session_id, timeout)

    def scroll(self, selector: str = None, x: int = None, y: int = None) -> Dict[str, Any]:
        """Scroll the page or an element. If a selector is provided, scrolls the element into view. If x and y coordinates are provided, scrolls the page to those absolute coordinates."""
        return self.client.browser_scroll(self.session_id, selector, x, y)

    def wait_for(self, condition: str, selector: str = None, timeout: int = 30000) -> Dict[str, Any]:
        """Wait for a condition"""
        return self.client.browser_wait(condition, self.session_id, selector, timeout)

    def screenshot(self, full_page: bool = True, clip: Dict = None, format: str = "png", quality: int = None) -> str:
        """Take a screenshot (returns base64 string)"""
        return self.client.browser_screenshot(self.session_id, full_page, clip, format, quality)

    def extract_elements(self, selector: str = None, attributes: List[str] = None) -> List[Dict[str, Any]]:
        """Extract DOM elements"""
        return self.client.browser_extract_elements(self.session_id, selector, attributes)

    def extract_content(self, url: str = None, include_interactive: bool = True,
                       include_anchors: bool = True, max_anchors: int = 50) -> Dict[str, Any]:
        """Extract LLM-friendly content with clean text, interactive elements, and anchors"""
        return self.client.browser_extract_content(self.session_id, url, include_interactive,
                                                   include_anchors, max_anchors)

    def close(self) -> bool:
        """Close this browser session"""
        if self._active:
            result = self.client.close_browser_session(self.session_id)
            self._active = False
            return result
        return True

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class BrowserManager:
    """Manager for browser automation functionality"""

    def __init__(self, client: 'InstaVM'):
        self.client = client
        self._default_session_id = None

    def create_session(self, viewport_width: int = 1920, viewport_height: int = 1080,
                      user_agent: str = None) -> BrowserSession:
        """Create a new browser session"""
        session_id = self.client.create_browser_session(viewport_width, viewport_height, user_agent)
        return BrowserSession(session_id, self.client)

    def _ensure_default_session(self) -> str:
        """Ensure a default browser session exists"""
        if not self._default_session_id:
            session = self.create_session()
            self._default_session_id = session.session_id
        return self._default_session_id

    def navigate(self, url: str, session_id: str = None, wait_timeout: int = 30000) -> Dict[str, Any]:
        """Navigate to URL (auto-creates session if none provided)"""
        if not session_id:
            session_id = self._ensure_default_session()
        return self.client.browser_navigate(url, session_id, wait_timeout)

    def click(self, selector: str, session_id: str = None, force: bool = False, timeout: int = 30000) -> Dict[str, Any]:
        """Click element by CSS selector"""
        if not session_id:
            session_id = self._ensure_default_session()
        return self.client.browser_click(selector, session_id, force, timeout)

    def type(self, selector: str, text: str, session_id: str = None, delay: int = 100, timeout: int = 30000) -> Dict[str, Any]:
        """Type text into element"""
        if not session_id:
            session_id = self._ensure_default_session()
        return self.client.browser_type(selector, text, session_id, delay, timeout)

    def fill(self, selector: str, value: str, session_id: str = None, timeout: int = 30000) -> Dict[str, Any]:
        """Fill form field"""
        if not session_id:
            session_id = self._ensure_default_session()
        return self.client.browser_fill(selector, value, session_id, timeout)

    def wait_for(self, condition: str, selector: str = None, session_id: str = None, timeout: int = 30000) -> Dict[str, Any]:
        """Wait for condition"""
        if not session_id:
            session_id = self._ensure_default_session()
        return self.client.browser_wait(condition, session_id, selector, timeout)

    def screenshot(self, session_id: str = None, full_page: bool = True, clip: Dict = None,
                  format: str = "png", quality: int = None) -> str:
        """Take screenshot (returns base64)"""
        if not session_id:
            session_id = self._ensure_default_session()
        return self.client.browser_screenshot(session_id, full_page, clip, format, quality)

    def extract_elements(self, selector: str = None, session_id: str = None, attributes: List[str] = None) -> List[Dict]:
        """Extract DOM elements"""
        if not session_id:
            session_id = self._ensure_default_session()
        return self.client.browser_extract_elements(session_id, selector, attributes)

    def extract_content(self, session_id: str = None, url: str = None,
                       include_interactive: bool = True, include_anchors: bool = True,
                       max_anchors: int = 50) -> Dict[str, Any]:
        """Extract LLM-friendly content (auto-creates session if none provided in cloud mode)"""
        if not self.client.local and not session_id:
            session_id = self._ensure_default_session()
        return self.client.browser_extract_content(session_id, url, include_interactive,
                                                   include_anchors, max_anchors)

    def close(self):
        """Closes the default browser session, if one was created."""
        if self._default_session_id:
            self.client.close_browser_session(self._default_session_id)
            self._default_session_id = None


class InstaVM:
    def __init__(self, api_key=None, base_url="https://api.instavm.io", timeout=300, max_retries=0,
                 local=False, local_url="http://coderunner.local:8222",
                 cpu_count: Optional[int] = None, memory_mb: Optional[int] = None,
                 env: Optional[Dict[str, Any]] = None, metadata: Optional[Dict[str, Any]] = None):
        """
        Initialize InstaVM client

        Args:
            api_key: API key for cloud mode (not required for local mode)
            base_url: Base URL for cloud API (ignored if local=True)
            timeout: VM lifetime in seconds (used when creating session). Also used as default HTTP request timeout. Range: 20-86400 seconds. Default: 300
            max_retries: Maximum number of retries for failed requests
            local: Use local container instead of cloud API
            local_url: Local container URL (default: http://coderunner.local:8222)
            cpu_count: Optional vCPU count for new session-backed VM (omitted if None so API defaults apply)
            memory_mb: Optional memory allocation (MB) for new session-backed VM (omitted if None so API defaults apply)
            env: Optional environment variables for new session-backed VM
            metadata: Optional metadata labels for new session-backed VM
        """
        self.local = local
        self.base_url = local_url if local else base_url
        self.api_key = api_key
        self.session_id = None  # Code execution session
        self.timeout = timeout
        self.max_retries = max_retries
        self.cpu_count = cpu_count
        self.memory_mb = memory_mb
        self.env = env or {}
        self.metadata = metadata or {}
        self.session = requests.Session()

        # Browser automation manager
        self.browser = BrowserManager(self)
        self.vms = VMsManager(self)
        self.snapshots = SnapshotsManager(self)
        self.shares = SharesManager(self)
        self.custom_domains = CustomDomainsManager(self)
        self.computer_use = ComputerUseManager(self)
        self.api_keys = APIKeysManager(self)
        self.audit = AuditManager(self)
        self.webhooks = WebhooksManager(self)
        self.volumes = VolumesManager(self)

        # Only start cloud session if not in local mode
        if not self.local and self.api_key:
            self.start_session()

    def _ensure_not_local(self, operation_name: str):
        """Raise UnsupportedOperationError if in local mode"""
        if self.local:
            raise UnsupportedOperationError(
                f"{operation_name} is not supported in local mode. "
                f"This operation is only available when using the cloud API."
            )

    def _make_request(self, method: str, url: str, timeout: Optional[int] = None, **kwargs) -> requests.Response:
        """Make HTTP request with retry logic and proper error handling"""
        last_exception = None
        request_timeout = timeout if timeout is not None else self.timeout

        for attempt in range(self.max_retries + 1):
            try:
                response = self.session.request(
                    method, url, timeout=request_timeout, **kwargs
                )

                # Handle specific HTTP status codes
                if response.status_code == 401:
                    raise AuthenticationError("Invalid API key or session expired")
                elif response.status_code == 429:
                    detail = response.json().get('detail', 'Rate limit exceeded')
                    raise RateLimitError(detail)
                elif response.status_code >= 500:
                    # More details for debugging
                    status_reason = response.reason or ""
                    try:
                        error_text = response.text.strip()
                    except Exception:
                        error_text = "<No response body>"

                    if response.status_code == 504:
                        raise NetworkError(
                            f"504 Gateway Timeout: The server (or a proxy) didn't get a timely response from the upstream service.\n"
                            f"Reason: {status_reason}\n"
                            f"Details: {error_text}"
                        )
                    else:
                        raise NetworkError(
                            f"Server error: {response.status_code} {status_reason}\n"
                            f"Details: {error_text}"
                        )


                response.raise_for_status()
                return response

            except requests.exceptions.Timeout as e:
                last_exception = NetworkError(f"Request timeout after {request_timeout}s")
            except requests.exceptions.ConnectionError as e:
                last_exception = NetworkError(f"Connection failed: {str(e)}")
            except (AuthenticationError, RateLimitError) as e:
                # Don't retry these
                raise e
            except requests.exceptions.HTTPError as e:
                if e.response.status_code < 500:
                    # Re-raise the original exception to allow for specific handling by the caller
                    raise
                last_exception = NetworkError(f"HTTP error: {str(e)}")

            if attempt < self.max_retries:
                # Exponential backoff
                time.sleep(2 ** attempt)

        raise last_exception or NetworkError("Max retries exceeded")

    def _ensure_api_key(self):
        """Raise AuthenticationError if API key is not set."""
        if not self.api_key:
            raise AuthenticationError("API key not set")

    def _auth_headers(self, extra: Optional[Dict[str, str]] = None) -> Dict[str, str]:
        """Build auth headers for cloud endpoints."""
        headers = {"X-API-Key": self.api_key} if self.api_key else {}
        if extra:
            headers.update(extra)
        return headers

    def _request_json(self, method: str, path: str, *,
                      params: Optional[Dict[str, Any]] = None,
                      json_data: Optional[Dict[str, Any]] = None,
                      data: Optional[Dict[str, Any]] = None,
                      headers: Optional[Dict[str, str]] = None,
                      timeout: Optional[int] = None,
                      require_auth: bool = True) -> Any:
        """Issue an API request and parse JSON response when possible."""
        self._ensure_not_local(path)
        if require_auth:
            self._ensure_api_key()

        url = f"{self.base_url}{path}"
        request_headers = self._auth_headers(headers) if require_auth else (headers or {})
        response = self._make_request(
            method,
            url,
            params=params,
            json=json_data,
            data=data,
            headers=request_headers,
            timeout=timeout
        )
        if not response.content:
            return {"success": True}
        try:
            return response.json()
        except ValueError:
            return response.text

    def start_session(self):
        self._ensure_not_local("Session management")

        if not self.api_key:
            raise AuthenticationError("API key not set. Please provide an API key or create one first.")

        url = f"{self.base_url}/v1/sessions/session"
        data = {
            "api_key": self.api_key,
            "vm_lifetime_seconds": self.timeout
        }
        if self.memory_mb is not None:
            data["memory_mb"] = self.memory_mb
        if self.cpu_count is not None:
            data["vcpu_count"] = self.cpu_count
        if self.metadata:
            data["metadata"] = self.metadata
        if self.env:
            data["env"] = self.env

        try:
            response = self._make_request("POST", url, json=data)
            result = response.json()
            self.session_id = result.get("session_id")
            if not self.session_id:
                raise SessionError("Failed to get session ID from server response")
            return self.session_id
        except Exception as e:
            if isinstance(e, (InstaVMError)):
                raise e
            raise SessionError(f"Failed to start session: {str(e)}")

    def execute(self, command: str, language: Optional[str] = None, timeout: Optional[int] = None) -> Dict[str, Any]:
        """
        Execute a command in the VM

        Args:
            command: Command to execute
            language: Programming language (optional)
            timeout: Request timeout in seconds (used both for HTTP request timeout and sent to API)

        Returns:
            Dict containing execution results
        """
        # In local mode, session_id is not required
        if not self.local and not self.session_id:
            raise SessionError("Session ID not set. Please start a session first.")

        url = f"{self.base_url}/execute"
        headers = {"X-API-Key": self.api_key} if self.api_key else {}
        data = {"command": command}

        # Only include session_id in cloud mode
        if not self.local:
            data["session_id"] = self.session_id
        if language:
            data["language"] = language
        if timeout is not None:
            data["timeout"] = timeout

        # Use custom timeout for this request, or fall back to instance default
        request_timeout = timeout if timeout is not None else self.timeout

        try:
            response = self._make_request("POST", url, headers=headers, json=data, timeout=request_timeout)
            return response.json()
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise e
            raise ExecutionError(f"Failed to execute command: {str(e)}")

    def get_usage(self) -> Dict[str, Any]:
        self._ensure_not_local("Usage tracking")

        if not self.session_id:
            raise SessionError("Session ID not set. Please start a session first.")

        safe_session_id = _encode_path_segment(self.session_id)
        url = f"{self.base_url}/v1/sessions/usage/{safe_session_id}"

        try:
            response = self._make_request("GET", url)
            return response.json()
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise e
            raise InstaVMError(f"Failed to get usage: {str(e)}")

    def get_session_info(self, session_id: Optional[str] = None) -> Dict[str, Any]:
        """Get session record details from /v1/sessions/session/{session_id}."""
        self._ensure_not_local("Session management")
        self._ensure_api_key()

        sid = session_id or self.session_id
        if not sid:
            raise SessionError("Session ID not set. Please start a session first.")

        try:
            safe_sid = _encode_path_segment(sid)
            return self._request_json("GET", f"/v1/sessions/session/{safe_sid}")
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise SessionError(f"Failed to get session info: {str(e)}")

    def get_session_app_url(self, session_id: Optional[str] = None,
                            port: Optional[int] = None) -> Dict[str, Any]:
        """Get app URL details for a session.

        Args:
            session_id: Session ID (uses current session if not provided)
            port: Port to expose (1-65535, default: 80)
        """
        self._ensure_not_local("Session management")
        self._ensure_api_key()

        sid = session_id or self.session_id
        if not sid:
            raise SessionError("Session ID not set. Please start a session first.")

        try:
            safe_sid = _encode_path_segment(sid)
            params: Optional[Dict[str, Any]] = None
            if port is not None:
                params = {"port": port}
            return self._request_json("GET", f"/v1/sessions/app-url/{safe_sid}",
                                      params=params)
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise SessionError(f"Failed to get session app URL: {str(e)}")

    def list_sandboxes(self, metadata: Optional[Dict[str, Any]] = None,
                       limit: Optional[int] = None) -> List[Dict[str, Any]]:
        """List sandbox records from /v1/sessions/sandboxes.

        Args:
            metadata: Optional metadata filter (dict serialized as JSON query param)
            limit: Maximum number of results (1-1000, default: 100)
        """
        self._ensure_not_local("Session management")
        self._ensure_api_key()

        try:
            params: Dict[str, Any] = {}
            if metadata is not None:
                params["metadata"] = json.dumps(metadata)
            if limit is not None:
                params["limit"] = limit
            result = self._request_json("GET", "/v1/sessions/sandboxes",
                                        params=params or None)
            return result if isinstance(result, list) else []
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise SessionError(f"Failed to list sandboxes: {str(e)}")

    def kill(self, session_id: Optional[str] = None) -> Dict[str, Any]:
        """Kill VM associated with a session via /kill."""
        self._ensure_not_local("VM kill")
        self._ensure_api_key()

        sid = session_id or self.session_id
        if not sid:
            raise SessionError("Session ID not set. Please start a session first.")

        try:
            result = self._request_json("POST", "/kill", json_data={"session_id": sid})
            if not session_id and sid == self.session_id:
                self.session_id = None
            return result
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise SessionError(f"Failed to kill VM: {str(e)}")

    def upload_file(self, file_path: str, remote_path: str,
                    recursive: bool = False) -> Dict[str, Any]:
        self._ensure_not_local("File upload")

        if not self.session_id:
            raise SessionError("Session ID not set. Please start a session first.")

        url = f"{self.base_url}/upload"
        headers = {"X-API-Key": self.api_key}
        try:
            with open(file_path, 'rb') as f:
                files = {'file': (file_path, f)}
                data = {
                    "remote_path": remote_path,
                    "session_id": self.session_id,
                    "recursive": str(recursive).lower()  # ensures 'true'/'false'
                }
                response = self._make_request("POST", url, headers=headers, data=data, files=files)
            return response.json()
        except FileNotFoundError:
            raise InstaVMError(f"File not found: {file_path}")
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise InstaVMError(f"Failed to upload file: {str(e)}")

    def download_file(self, filename: str, local_path: Optional[str] = None) -> Dict[str, Any]:
        """Download a file from the remote VM

        Args:
            filename: Name of the file to download from the remote VM
            local_path: Optional local path to save the file. If not provided, returns file content in response

        Returns:
            Dict containing download status and file information
        """
        self._ensure_not_local("File download")

        if not self.session_id:
            raise SessionError("Session ID not set. Please start a session first.")

        url = f"{self.base_url}/download"
        headers = {"X-API-Key": self.api_key}
        data = {
            "filename": filename,
            "session_id": self.session_id
        }

        try:
            response = self._make_request("POST", url, headers=headers, data=data)

            # Parse JSON response
            response_data = response.json()

            # Extract and decode file content (assuming base64-encoded)
            encoded_content = response_data.get("content", "")
            if encoded_content:
                file_content = base64.b64decode(encoded_content)
            else:
                # Fallback: if no 'content' field, try raw response
                file_content = response.content

            if local_path:
                # Save to local file
                with open(local_path, 'wb') as f:
                    f.write(file_content)
                return {
                    "status": "success",
                    "filename": filename,
                    "local_path": local_path,
                    "size": len(file_content)
                }
            else:
                # Return content in response
                return {
                    "status": "success",
                    "filename": filename,
                    "content": file_content,
                    "size": len(file_content)
                }
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise InstaVMError(f"Failed to download file: {str(e)}")

    def execute_async(self, command: str, language: Optional[str] = None, timeout: Optional[int] = None) -> Dict[str, Any]:
        """Execute command asynchronously"""
        self._ensure_not_local("Async execution")

        if not self.session_id:
            raise SessionError("Session ID not set. Please start a session first.")

        url = f"{self.base_url}/execute_async"
        headers = {"X-API-Key": self.api_key}
        data = {
            "command": command,
            "session_id": self.session_id,
        }

        if language:
            data["language"] = language
        if timeout is not None:
            data["timeout"] = timeout

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            return response.json()
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise e
            raise ExecutionError(f"Failed to execute command asynchronously: {str(e)}")

    def get_task_result(self, task_id: str, poll_interval: int = 1, timeout: int = 60) -> Dict[str, Any]:
        """Poll for async task result.

        Args:
            task_id: The task ID from execute_async
            poll_interval: Seconds between polls (default: 1)
            timeout: Maximum seconds to wait (default: 60)

        Returns:
            Dict containing:
                - stdout: Task stdout output
                - stderr: Task stderr output
                - execution_time: Task execution time (if available)
                - cpu_time: Task CPU time (if available)

        Raises:
            TimeoutError: If task doesn't complete within timeout
            ExecutionError: If task retrieval fails
        """
        self._ensure_not_local("Async task result retrieval")

        start_time = time.time()

        while time.time() - start_time < timeout:
            try:
                safe_task_id = _encode_path_segment(task_id)
                url = f"{self.base_url}/v1/executions/{safe_task_id}"
                headers = {"X-API-Key": self.api_key}

                response = self._make_request("GET", url, headers=headers)
                task = response.json()

                if task.get("is_complete"):
                    # Output is already decoded as string (not base64 despite field name)
                    stdout = task.get("serialized_stdout", "")
                    stderr = task.get("serialized_stderr", "")

                    return {
                        "stdout": stdout,
                        "stderr": stderr,
                        "cpu_time": task.get("cpu_time"),
                        "created_at": task.get("created_at"),
                        "execution_time": task.get("execution_time")
                    }

                time.sleep(poll_interval)

            except (AuthenticationError, RateLimitError, NetworkError) as e:
                # Don't retry these specific errors
                raise e
            except Exception as e:
                # Continue polling on transient errors
                if time.time() - start_time >= timeout:
                    raise ExecutionError(f"Failed to get task result: {str(e)}")
                time.sleep(poll_interval)

        raise TimeoutError(f"Task {task_id} timed out after {timeout}s")

    def execute_streaming(self, command: str, on_output: Optional[Callable[[str], None]] = None) -> Generator[str, None, None]:
        """Execute command with real-time streaming output (deprecated - streaming not supported by current API)"""
        import warnings
        warnings.warn("execute_streaming is deprecated. The API does not support streaming execution. Use execute() or execute_async() instead.", DeprecationWarning, stacklevel=2)

        # Fallback to regular execution
        result = self.execute(command)
        output = result.get('output', str(result))
        if on_output:
            on_output(output)
        yield output

    # --- Egress Policy Methods ---

    def set_session_egress(self, session_id: str = None, *,
                           allow_package_managers: bool = True,
                           allow_http: bool = True,
                           allow_https: bool = True,
                           allowed_domains: List[str] = None,
                           allowed_cidrs: List[str] = None) -> dict:
        """Set egress policy for a session. Uses current session if session_id omitted."""
        self._ensure_not_local("Egress policy management")
        self._ensure_api_key()

        sid = session_id or self.session_id
        if not sid:
            raise SessionError("Session ID not set. Please start a session first.")

        safe_sid = _encode_path_segment(sid)
        url = f"{self.base_url}/v1/egress/session/{safe_sid}"
        headers = {"X-API-Key": self.api_key}
        data = {
            "allow_public_repos": allow_package_managers,
            "allow_http": allow_http,
            "allow_https": allow_https,
            "allowed_domains": allowed_domains or [],
            "allowed_cidrs": allowed_cidrs or [],
        }

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            return response.json()
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise InstaVMError(f"Failed to set session egress policy: {str(e)}")

    def get_session_egress(self, session_id: str = None) -> dict:
        """Get egress policy for a session. Uses current session if session_id omitted."""
        self._ensure_not_local("Egress policy management")
        self._ensure_api_key()

        sid = session_id or self.session_id
        if not sid:
            raise SessionError("Session ID not set. Please start a session first.")

        safe_sid = _encode_path_segment(sid)
        url = f"{self.base_url}/v1/egress/session/{safe_sid}"
        headers = {"X-API-Key": self.api_key}

        try:
            response = self._make_request("GET", url, headers=headers)
            return response.json()
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise InstaVMError(f"Failed to get session egress policy: {str(e)}")

    def set_vm_egress(self, vm_id: str, *,
                      allow_package_managers: bool = True,
                      allow_http: bool = True,
                      allow_https: bool = True,
                      allowed_domains: List[str] = None,
                      allowed_cidrs: List[str] = None) -> dict:
        """Set egress policy for a specific VM."""
        self._ensure_not_local("Egress policy management")
        self._ensure_api_key()

        safe_vm_id = _encode_path_segment(vm_id)
        url = f"{self.base_url}/v1/egress/vm/{safe_vm_id}"
        headers = {"X-API-Key": self.api_key}
        data = {
            "allow_public_repos": allow_package_managers,
            "allow_http": allow_http,
            "allow_https": allow_https,
            "allowed_domains": allowed_domains or [],
            "allowed_cidrs": allowed_cidrs or [],
        }

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            return response.json()
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise InstaVMError(f"Failed to set VM egress policy: {str(e)}")

    def get_vm_egress(self, vm_id: str) -> dict:
        """Get egress policy for a specific VM."""
        self._ensure_not_local("Egress policy management")
        self._ensure_api_key()

        safe_vm_id = _encode_path_segment(vm_id)
        url = f"{self.base_url}/v1/egress/vm/{safe_vm_id}"
        headers = {"X-API-Key": self.api_key}

        try:
            response = self._make_request("GET", url, headers=headers)
            return response.json()
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise InstaVMError(f"Failed to get VM egress policy: {str(e)}")

    # --- SSH Key Methods ---

    def add_ssh_key(self, public_key: str) -> dict:
        """Add an SSH public key. Returns the key record with fingerprint."""
        self._ensure_not_local("SSH key management")
        self._ensure_api_key()

        url = f"{self.base_url}/v1/ssh-keys"
        headers = {"X-API-Key": self.api_key}
        data = {"public_key": public_key}

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            return response.json()
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise InstaVMError(f"Failed to add SSH key: {str(e)}")

    def list_ssh_keys(self) -> List[dict]:
        """List all active SSH keys. Returns list of key records."""
        self._ensure_not_local("SSH key management")
        self._ensure_api_key()

        url = f"{self.base_url}/v1/ssh-keys"
        headers = {"X-API-Key": self.api_key}

        try:
            response = self._make_request("GET", url, headers=headers)
            result = response.json()
            return result.get("keys", [])
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise InstaVMError(f"Failed to list SSH keys: {str(e)}")

    def delete_ssh_key(self, key_id: int) -> dict:
        """Delete an SSH key by ID."""
        self._ensure_not_local("SSH key management")
        self._ensure_api_key()

        safe_key_id = _encode_path_segment(key_id)
        url = f"{self.base_url}/v1/ssh-keys/{safe_key_id}"
        headers = {"X-API-Key": self.api_key}

        try:
            response = self._make_request("DELETE", url, headers=headers)
            return response.json()
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise InstaVMError(f"Failed to delete SSH key: {str(e)}")

    def close_session(self) -> bool:
        """Close the current session (sessions auto-expire on server side)"""
        if not self.session_id:
            return True

        # Note: API doesn't provide explicit session deletion endpoint
        # Sessions will auto-expire on the server side
        print(f"Info: Session {self.session_id} will auto-expire on server side.")
        self.session_id = None
        return True

    def is_session_active(self) -> bool:
        """Check if current session is still active by checking VM status"""
        if not self.session_id:
            return False

        try:
            safe_session_id = _encode_path_segment(self.session_id)
            url = f"{self.base_url}/v1/sessions/status/{safe_session_id}"
            headers = {"X-API-Key": self.api_key} if self.api_key else {}
            response = self._make_request("GET", url, headers=headers)
            result = response.json()
            return result.get("vm_status") == "active"
        except (SessionError, AuthenticationError, NetworkError):
            return False
        except Exception:
            return False

    def __enter__(self):
        """Context manager entry"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - automatically close session"""
        self.close_session()
        # Parameters are standard context manager signature

    # Browser Automation Methods
    def create_browser_session(self, viewport_width: int = 1920, viewport_height: int = 1080,
                              user_agent: str = None) -> str:
        """Create a new browser session"""
        self._ensure_not_local("Browser session management")
        self._ensure_api_key()

        url = f"{self.base_url}/v1/browser/sessions/"
        headers = {"X-API-Key": self.api_key}
        data = {
            "viewport_width": viewport_width,
            "viewport_height": viewport_height
        }
        if user_agent:
            data["user_agent"] = user_agent

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            result = response.json()
            return result.get("session_id")
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 429:
                raise QuotaExceededError("Browser API quota exceeded")
            raise BrowserSessionError(f"Failed to create browser session: {str(e)}")
        except Exception as e:
            raise BrowserSessionError(f"Failed to create browser session: {str(e)}")

    def get_browser_session(self, session_id: str) -> Dict[str, Any]:
        """Get browser session information"""
        self._ensure_not_local("Browser session management")
        self._ensure_api_key()

        safe_session_id = _encode_path_segment(session_id)
        url = f"{self.base_url}/v1/browser/sessions/{safe_session_id}"
        headers = {"X-API-Key": self.api_key}

        try:
            response = self._make_request("GET", url, headers=headers)
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 404:
                raise BrowserSessionError("Browser session not found or expired")
            raise BrowserSessionError(f"Failed to get browser session: {str(e)}")
        except Exception as e:
            raise BrowserSessionError(f"Failed to get browser session: {str(e)}")

    def close_browser_session(self, session_id: str) -> bool:
        """Close a browser session"""
        self._ensure_not_local("Browser session management")
        self._ensure_api_key()

        safe_session_id = _encode_path_segment(session_id)
        url = f"{self.base_url}/v1/browser/sessions/{safe_session_id}"
        headers = {"X-API-Key": self.api_key}

        try:
            response = self._make_request("DELETE", url, headers=headers)
            return response.status_code == 200
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 404:
                return True  # Already closed
            raise BrowserSessionError(f"Failed to close browser session: {str(e)}")
        except Exception as e:
            raise BrowserSessionError(f"Failed to close browser session: {str(e)}")

    def list_browser_sessions(self) -> List[Dict[str, Any]]:
        """List active browser sessions"""
        self._ensure_not_local("Browser session management")
        self._ensure_api_key()

        url = f"{self.base_url}/v1/browser/sessions/"
        headers = {"X-API-Key": self.api_key}

        try:
            response = self._make_request("GET", url, headers=headers)
            result = response.json()
            return result.get("sessions", [])
        except Exception as e:
            raise BrowserSessionError(f"Failed to list browser sessions: {str(e)}")

    def browser_navigate(self, url: str, session_id: str = None, wait_timeout: int = 30000) -> Dict[str, Any]:
        """Navigate to a URL in browser session

        Args:
            url: URL to navigate to
            session_id: Browser session ID (not required in local mode)
            wait_timeout: Timeout in milliseconds
        """
        # In cloud mode, require API key and session_id
        if not self.local:
            self._ensure_api_key()
            if not session_id:
                raise ValueError("session_id is required in cloud mode")

        api_url = f"{self.base_url}/v1/browser/interactions/navigate"
        headers = {"X-API-Key": self.api_key} if self.api_key else {}
        data = {"url": url, "wait_timeout": wait_timeout}

        # Only include session_id in cloud mode
        if not self.local and session_id:
            data["session_id"] = session_id

        try:
            response = self._make_request("POST", api_url, headers=headers, json=data)
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 404:
                raise BrowserSessionError("Browser session not found or expired")
            if e.response.status_code == 408:
                raise BrowserTimeoutError("Navigation timeout")
            raise BrowserNavigationError(f"Failed to navigate: {str(e)}")
        except Exception as e:
            raise BrowserNavigationError(f"Failed to navigate: {str(e)}")

    def browser_click(self, selector: str, session_id: str, force: bool = False, timeout: int = 30000) -> Dict[str, Any]:
        """Click an element in browser session"""
        self._ensure_not_local("Browser click interaction")
        self._ensure_api_key()

        url = f"{self.base_url}/v1/browser/interactions/click"
        headers = {"X-API-Key": self.api_key}
        data = {
            "selector": selector,
            "session_id": session_id,
            "force": force,
            "timeout": timeout
        }

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 404:
                raise ElementNotFoundError(f"Element not found: {selector}")
            if e.response.status_code == 408:
                raise BrowserTimeoutError("Click timeout")
            raise BrowserInteractionError(f"Failed to click element: {str(e)}")
        except Exception as e:
            raise BrowserInteractionError(f"Failed to click element: {str(e)}")

    def browser_type(self, selector: str, text: str, session_id: str, delay: int = 100, timeout: int = 30000) -> Dict[str, Any]:
        """Type text into an element"""
        self._ensure_not_local("Browser type interaction")
        self._ensure_api_key()

        url = f"{self.base_url}/v1/browser/interactions/type"
        headers = {"X-API-Key": self.api_key}
        data = {
            "selector": selector,
            "text": text,
            "session_id": session_id,
            "delay": delay,
            "timeout": timeout
        }

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 404:
                raise ElementNotFoundError(f"Element not found: {selector}")
            if e.response.status_code == 408:
                raise BrowserTimeoutError("Type timeout")
            raise BrowserInteractionError(f"Failed to type text: {str(e)}")
        except Exception as e:
            raise BrowserInteractionError(f"Failed to type text: {str(e)}")

    def browser_fill(self, selector: str, value: str, session_id: str, timeout: int = 30000) -> Dict[str, Any]:
        """Fill a form field"""
        self._ensure_not_local("Browser fill interaction")
        self._ensure_api_key()

        url = f"{self.base_url}/v1/browser/interactions/fill"
        headers = {"X-API-Key": self.api_key}
        data = {
            "selector": selector,
            "value": value,
            "session_id": session_id,
            "timeout": timeout
        }

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 404:
                raise ElementNotFoundError(f"Element not found: {selector}")
            if e.response.status_code == 408:
                raise BrowserTimeoutError("Fill timeout")
            raise BrowserInteractionError(f"Failed to fill form: {str(e)}")
        except Exception as e:
            raise BrowserInteractionError(f"Failed to fill form: {str(e)}")

    def browser_scroll(self, session_id: str, selector: str = None, x: int = None, y: int = None) -> Dict[str, Any]:
        """Scroll the page or an element"""
        self._ensure_not_local("Browser scroll interaction")
        self._ensure_api_key()

        if selector is None and x is None and y is None:
            raise ValueError("At least one of 'selector', 'x', or 'y' must be provided for scrolling.")

        url = f"{self.base_url}/v1/browser/interactions/scroll"
        headers = {"X-API-Key": self.api_key}
        data = {"session_id": session_id}
        if selector:
            data["selector"] = selector
        if x is not None:
            data["x"] = x
        if y is not None:
            data["y"] = y

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            return response.json()
        except Exception as e:
            raise BrowserInteractionError(f"Failed to scroll: {str(e)}")

    def browser_wait(self, condition: str, session_id: str, selector: str = None, timeout: int = 30000) -> Dict[str, Any]:
        """Wait for a condition to be met"""
        self._ensure_not_local("Browser wait interaction")
        self._ensure_api_key()

        url = f"{self.base_url}/v1/browser/interactions/wait"
        headers = {"X-API-Key": self.api_key}
        data = {
            "condition": condition,
            "session_id": session_id,
            "timeout": timeout
        }
        if selector:
            data["selector"] = selector

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 408:
                raise BrowserTimeoutError(f"Wait timeout for condition: {condition}")
            raise BrowserInteractionError(f"Failed to wait for condition: {str(e)}")
        except Exception as e:
            raise BrowserInteractionError(f"Failed to wait for condition: {str(e)}")

    def browser_screenshot(self, session_id: str, full_page: bool = True, clip: Dict = None,
                          format: str = "png", quality: int = None) -> str:
        """Take a screenshot (returns base64 string)"""
        self._ensure_not_local("Browser screenshot")
        self._ensure_api_key()

        url = f"{self.base_url}/v1/browser/interactions/screenshot"
        headers = {"X-API-Key": self.api_key}
        data = {
            "session_id": session_id,
            "full_page": full_page,
            "format": format
        }
        if clip:
            data["clip"] = clip
        if quality:
            data["quality"] = quality

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            result = response.json()
            return result.get("screenshot", "")
        except Exception as e:
            raise BrowserInteractionError(f"Failed to take screenshot: {str(e)}")

    def browser_extract_elements(self, session_id: str, selector: str = None, attributes: List[str] = None) -> List[Dict[str, Any]]:
        """Extract DOM elements"""
        self._ensure_not_local("Browser element extraction")
        self._ensure_api_key()

        url = f"{self.base_url}/v1/browser/interactions/extract"
        headers = {"X-API-Key": self.api_key}
        data = {"session_id": session_id}
        if selector:
            data["selector"] = selector
        if attributes:
            data["attributes"] = attributes

        try:
            response = self._make_request("POST", url, headers=headers, json=data)
            result = response.json()
            return result.get("elements", [])
        except Exception as e:
            raise BrowserInteractionError(f"Failed to extract elements: {str(e)}")

    def browser_extract_content(self, session_id: str = None, url: str = None,
                                include_interactive: bool = True,
                                include_anchors: bool = True, max_anchors: int = 50) -> Dict[str, Any]:
        """
        Extract LLM-friendly content from current page.

        Returns clean article content, interactive elements, and content anchors
        for intelligent browser automation.

        Args:
            session_id: Browser session ID (not required in local mode, required in cloud mode)
            url: URL to extract content from (required in local mode, optional in cloud mode)
            include_interactive: Include interactive elements mapping (default: True)
            include_anchors: Include content-to-selector anchors (default: True)
            max_anchors: Maximum number of content anchors (default: 50)

        Returns:
            Dict with:
            - readable_content: Clean article text (no JS/CSS/ads)
            - interactive_elements: Clickable/typeable elements with selectors
            - content_anchors: Text snippets mapped to DOM selectors

        Example:
            ```python
            # Cloud mode
            content = client.browser_extract_content(session_id="session123")

            # Local mode
            content = client.browser_extract_content(url="https://example.com")

            # LLM reads clean content
            article = content['readable_content']['content']

            # LLM finds "Sign Up" in content
            # LLM searches content_anchors for selector
            for anchor in content['content_anchors']:
                if 'sign up' in anchor['text'].lower():
                    selector = anchor['selector']
                    client.browser_click(selector, session_id)
                    break
            ```
        """
        # In cloud mode, require API key and session_id
        if not self.local:
            self._ensure_api_key()
            if not session_id:
                raise ValueError("session_id is required in cloud mode")
        else:
            # In local mode, require url
            if not url:
                raise ValueError("url is required in local mode")

        api_url = f"{self.base_url}/v1/browser/interactions/content"
        headers = {"X-API-Key": self.api_key} if self.api_key else {}
        data = {
            "include_interactive": include_interactive,
            "include_anchors": include_anchors,
            "max_anchors": max_anchors
        }

        # Include session_id in cloud mode
        if not self.local and session_id:
            data["session_id"] = session_id

        # Include url in local mode
        if self.local and url:
            data["url"] = url

        try:
            response = self._make_request("POST", api_url, headers=headers, json=data)
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 404:
                raise BrowserSessionError("Browser session not found or expired")
            raise BrowserInteractionError(f"Failed to extract content: {str(e)}")
        except json.JSONDecodeError as e:
            raise BrowserInteractionError(f"Failed to parse server response: {str(e)}")


class VMsManager:
    """VM lifecycle API wrapper."""

    def __init__(self, client: InstaVM):
        self.client = client

    def create(self, wait: bool = True,
               volumes: Optional[List[Dict[str, Any]]] = None,
               **payload: Any) -> Dict[str, Any]:
        body = dict(payload)
        if volumes is not None:
            body["volumes"] = volumes
        return self.client._request_json(
            "POST",
            "/v1/vms",
            params={"wait": str(wait).lower()},
            json_data=body
        )

    def list(self) -> List[Dict[str, Any]]:
        """List VMs via GET /v1/vms."""
        result = self.client._request_json("GET", "/v1/vms")
        return result if isinstance(result, list) else []

    def list_all_records(self) -> List[Dict[str, Any]]:
        """List all VM records via GET /v1/vms/."""
        result = self.client._request_json("GET", "/v1/vms/")
        return result if isinstance(result, list) else []

    def get(self, vm_id: str) -> Dict[str, Any]:
        safe_id = _encode_path_segment(vm_id)
        return self.client._request_json("GET", f"/v1/vms/{safe_id}")

    def update(self, vm_id: str, **payload: Any) -> Dict[str, Any]:
        safe_vm_id = _encode_path_segment(vm_id)
        return self.client._request_json(
            "PATCH",
            f"/v1/vms/{safe_vm_id}",
            json_data=payload
        )

    def delete(self, vm_id: str) -> Dict[str, Any]:
        safe_vm_id = _encode_path_segment(vm_id)
        return self.client._request_json("DELETE", f"/v1/vms/{safe_vm_id}")

    def clone(self, vm_id: str, wait: bool = True,
              snapshot_name: Optional[str] = None, **payload: Any) -> Dict[str, Any]:
        body = dict(payload)
        if snapshot_name is not None:
            body["snapshot_name"] = snapshot_name
        safe_vm_id = _encode_path_segment(vm_id)
        return self.client._request_json(
            "POST",
            f"/v1/vms/{safe_vm_id}/clone",
            params={"wait": str(wait).lower()},
            json_data=body
        )

    def snapshot(self, vm_id: str, wait: bool = True,
                 name: Optional[str] = None, **payload: Any) -> Dict[str, Any]:
        body = dict(payload)
        if name is not None:
            body["name"] = name
        safe_vm_id = _encode_path_segment(vm_id)
        return self.client._request_json(
            "POST",
            f"/v1/vms/{safe_vm_id}/snapshot",
            params={"wait": str(wait).lower()},
            json_data=body
        )

    def mount_volume(self, vm_id: str, volume_id: str, mount_path: str,
                     mode: str = "rw", checkpoint_id: Optional[str] = None,
                     wait: bool = True) -> Dict[str, Any]:
        body: Dict[str, Any] = {
            "volume_id": volume_id,
            "mount_path": mount_path
        }
        if mode:
            body["mode"] = mode
        if checkpoint_id is not None:
            body["checkpoint_id"] = checkpoint_id

        safe_vm_id = _encode_path_segment(vm_id)
        return self.client._request_json(
            "POST",
            f"/v1/vms/{safe_vm_id}/volumes",
            params={"wait": str(wait).lower()},
            json_data=body
        )

    def list_volumes(self, vm_id: str) -> List[Dict[str, Any]]:
        safe_vm_id = _encode_path_segment(vm_id)
        result = self.client._request_json("GET", f"/v1/vms/{safe_vm_id}/volumes")
        return result if isinstance(result, list) else []

    def unmount_volume(self, vm_id: str, volume_id: str, mount_path: str,
                       wait: bool = True) -> Dict[str, Any]:
        safe_vm_id = _encode_path_segment(vm_id)
        safe_volume_id = _encode_path_segment(volume_id)
        return self.client._request_json(
            "DELETE",
            f"/v1/vms/{safe_vm_id}/volumes/{safe_volume_id}",
            params={
                "mount_path": mount_path,
                "wait": str(wait).lower()
            }
        )


class SnapshotsManager:
    """Snapshot API wrapper."""

    def __init__(self, client: InstaVM):
        self.client = client

    def create(self, oci_image: str, name: Optional[str] = None,
               vcpu_count: Optional[int] = None, memory_mb: Optional[int] = None,
               build_args: Optional[Dict[str, Any]] = None,
               snapshot_type: Optional[str] = None, **payload: Any) -> Dict[str, Any]:
        body = dict(payload)
        body["oci_image"] = oci_image
        if name is not None:
            body["name"] = name
        if vcpu_count is not None:
            body["vcpu_count"] = vcpu_count
        if memory_mb is not None:
            body["memory_mb"] = memory_mb
        if build_args is not None:
            body["build_args"] = build_args
        if snapshot_type is not None:
            body["type"] = snapshot_type
        return self.client._request_json("POST", "/v1/snapshots", json_data=body)

    def list(self, snapshot_type: Optional[str] = None, **query: Any) -> List[Dict[str, Any]]:
        if snapshot_type is None and "type" in query:
            snapshot_type = query.pop("type")
        params = dict(query)
        if snapshot_type is not None:
            params["type"] = snapshot_type
        result = self.client._request_json("GET", "/v1/snapshots", params=params)
        return result if isinstance(result, list) else []

    def get(self, snapshot_id: str) -> Dict[str, Any]:
        safe_snapshot_id = _encode_path_segment(snapshot_id)
        return self.client._request_json("GET", f"/v1/snapshots/{safe_snapshot_id}")

    def delete(self, snapshot_id: str) -> Dict[str, Any]:
        safe_snapshot_id = _encode_path_segment(snapshot_id)
        return self.client._request_json("DELETE", f"/v1/snapshots/{safe_snapshot_id}")


class VolumesManager:
    """Volume lifecycle + file + checkpoint API wrapper."""

    def __init__(self, client: InstaVM):
        self.client = client

    def create(self, name: str, quota_bytes: int, **payload: Any) -> Dict[str, Any]:
        body = dict(payload)
        body["name"] = name
        body["quota_bytes"] = quota_bytes
        return self.client._request_json("POST", "/v1/volumes", json_data=body)

    def list(self, refresh_usage: bool = False) -> List[Dict[str, Any]]:
        result = self.client._request_json(
            "GET",
            "/v1/volumes",
            params={"refresh_usage": str(refresh_usage).lower()}
        )
        return result if isinstance(result, list) else []

    def get(self, volume_id: str, refresh_usage: bool = False) -> Dict[str, Any]:
        safe_volume_id = _encode_path_segment(volume_id)
        return self.client._request_json(
            "GET",
            f"/v1/volumes/{safe_volume_id}",
            params={"refresh_usage": str(refresh_usage).lower()}
        )

    def update(self, volume_id: str, name: Optional[str] = None,
               quota_bytes: Optional[int] = None, **payload: Any) -> Dict[str, Any]:
        body = dict(payload)
        if name is not None:
            body["name"] = name
        if quota_bytes is not None:
            body["quota_bytes"] = quota_bytes
        if not body:
            raise ValueError("At least one update field must be provided.")

        safe_volume_id = _encode_path_segment(volume_id)
        return self.client._request_json(
            "PATCH",
            f"/v1/volumes/{safe_volume_id}",
            json_data=body
        )

    def delete(self, volume_id: str) -> Dict[str, Any]:
        safe_volume_id = _encode_path_segment(volume_id)
        return self.client._request_json("DELETE", f"/v1/volumes/{safe_volume_id}")

    def create_checkpoint(self, volume_id: str, name: Optional[str] = None,
                          **payload: Any) -> Dict[str, Any]:
        body = dict(payload)
        if name is not None:
            body["name"] = name
        safe_volume_id = _encode_path_segment(volume_id)
        return self.client._request_json(
            "POST",
            f"/v1/volumes/{safe_volume_id}/checkpoints",
            json_data=body
        )

    def list_checkpoints(self, volume_id: str) -> List[Dict[str, Any]]:
        safe_volume_id = _encode_path_segment(volume_id)
        result = self.client._request_json("GET", f"/v1/volumes/{safe_volume_id}/checkpoints")
        return result if isinstance(result, list) else []

    def delete_checkpoint(self, volume_id: str, checkpoint_id: str) -> Dict[str, Any]:
        safe_volume_id = _encode_path_segment(volume_id)
        safe_checkpoint_id = _encode_path_segment(checkpoint_id)
        return self.client._request_json(
            "DELETE",
            f"/v1/volumes/{safe_volume_id}/checkpoints/{safe_checkpoint_id}"
        )

    def upload_file(self, volume_id: str, file_path: str, path: str,
                    overwrite: bool = False) -> Dict[str, Any]:
        self.client._ensure_not_local("Volume file upload")
        self.client._ensure_api_key()

        safe_volume_id = _encode_path_segment(volume_id)
        url = f"{self.client.base_url}/v1/volumes/{safe_volume_id}/files/upload"
        headers = self.client._auth_headers()
        data = {
            "path": path,
            "overwrite": str(overwrite).lower()
        }

        try:
            with open(file_path, 'rb') as f:
                files = {"file": (file_path, f)}
                response = self.client._make_request(
                    "POST",
                    url,
                    headers=headers,
                    data=data,
                    files=files
                )
            if not response.content:
                return {"success": True}
            return response.json()
        except FileNotFoundError:
            raise InstaVMError(f"File not found: {file_path}")
        except Exception as e:
            if isinstance(e, InstaVMError):
                raise
            raise InstaVMError(f"Failed to upload volume file: {str(e)}")

    def download_file(self, volume_id: str, path: str,
                      local_path: Optional[str] = None,
                      decode_base64: bool = True) -> Dict[str, Any]:
        safe_volume_id = _encode_path_segment(volume_id)
        result = self.client._request_json(
            "POST",
            f"/v1/volumes/{safe_volume_id}/files/download",
            json_data={"path": path}
        )

        if local_path and isinstance(result, dict):
            encoded = result.get("content", "")
            raw_bytes: bytes
            if isinstance(encoded, str):
                if decode_base64:
                    try:
                        raw_bytes = base64.b64decode(encoded)
                    except (ValueError, binascii.Error):
                        raw_bytes = encoded.encode("utf-8")
                else:
                    raw_bytes = encoded.encode("utf-8")
            elif isinstance(encoded, bytes):
                raw_bytes = encoded
            else:
                raw_bytes = b""

            with open(local_path, 'wb') as f:
                f.write(raw_bytes)

            output = dict(result)
            output["local_path"] = local_path
            output["saved_bytes"] = len(raw_bytes)
            return output

        return result if isinstance(result, dict) else {"result": result}

    def list_files(self, volume_id: str, prefix: str = "",
                   recursive: bool = True, limit: int = 1000) -> List[Dict[str, Any]]:
        safe_volume_id = _encode_path_segment(volume_id)
        result = self.client._request_json(
            "GET",
            f"/v1/volumes/{safe_volume_id}/files",
            params={
                "prefix": prefix,
                "recursive": str(recursive).lower(),
                "limit": limit
            }
        )
        return result if isinstance(result, list) else []

    def delete_file(self, volume_id: str, path: str) -> Dict[str, Any]:
        safe_volume_id = _encode_path_segment(volume_id)
        return self.client._request_json(
            "DELETE",
            f"/v1/volumes/{safe_volume_id}/files",
            params={"path": path}
        )


class SharesManager:
    """Share API wrapper."""

    def __init__(self, client: InstaVM):
        self.client = client

    def create(self, port: int, session_id: Optional[str] = None,
               vm_id: Optional[str] = None, is_public: bool = False) -> Dict[str, Any]:
        body: Dict[str, Any] = {
            "port": port,
            "is_public": is_public
        }
        if session_id is not None:
            body["session_id"] = session_id
        elif self.client.session_id:
            body["session_id"] = self.client.session_id
        if vm_id is not None:
            body["vm_id"] = vm_id
        return self.client._request_json("POST", "/v1/shares", json_data=body)

    def update(self, share_id: str, is_public: Optional[bool] = None,
               revoke: Optional[bool] = None) -> Dict[str, Any]:
        body: Dict[str, Any] = {}
        if is_public is not None:
            body["is_public"] = is_public
        if revoke is not None:
            body["revoke"] = revoke
        if not body:
            raise ValueError("At least one of 'is_public' or 'revoke' must be provided.")
        safe_share_id = _encode_path_segment(share_id)
        return self.client._request_json(
            "PATCH",
            f"/v1/shares/{safe_share_id}",
            json_data=body
        )


class CustomDomainsManager:
    """Custom domain API wrapper."""

    def __init__(self, client: InstaVM):
        self.client = client

    def create(self, domain: str, share_id: str, dns_provider: str,
               dns_credentials: Optional[Dict[str, str]] = None) -> Dict[str, Any]:
        body: Dict[str, Any] = {
            "domain": domain,
            "share_id": share_id,
            "dns_provider": dns_provider,
        }
        if dns_credentials is not None:
            body["dns_credentials"] = dns_credentials
        return self.client._request_json("POST", "/v1/custom-domains", json_data=body)

    def list(self) -> List[Dict[str, Any]]:
        result = self.client._request_json("GET", "/v1/custom-domains")
        return result if isinstance(result, list) else []

    def health(self, domain_id: int) -> Dict[str, Any]:
        safe_domain_id = _encode_path_segment(domain_id)
        return self.client._request_json("GET", f"/v1/custom-domains/{safe_domain_id}/health")

    def verify(self, domain_id: int) -> Dict[str, Any]:
        safe_domain_id = _encode_path_segment(domain_id)
        return self.client._request_json("POST", f"/v1/custom-domains/{safe_domain_id}/verify")

    def delete(self, domain_id: int) -> Dict[str, Any]:
        safe_domain_id = _encode_path_segment(domain_id)
        return self.client._request_json("DELETE", f"/v1/custom-domains/{safe_domain_id}")


class ComputerUseManager:
    """Computer-use viewer + proxy API wrapper."""

    def __init__(self, client: InstaVM):
        self.client = client

    def viewer_url(self, session_id: str) -> Dict[str, Any]:
        safe_session_id = _encode_path_segment(session_id)
        return self.client._request_json("GET", f"/v1/computeruse/{safe_session_id}/viewer-url")

    def proxy(self, session_id: str, path: str, method: str = "GET",
              params: Optional[Dict[str, Any]] = None,
              json_data: Optional[Dict[str, Any]] = None,
              data: Optional[Dict[str, Any]] = None) -> Any:
        safe_session_id = _encode_path_segment(session_id)
        clean_path = _encode_proxy_path(path)
        return self.client._request_json(
            method.upper(),
            f"/v1/computeruse/{safe_session_id}/{clean_path}",
            params=params,
            json_data=json_data,
            data=data
        )

    def get(self, session_id: str, path: str, params: Optional[Dict[str, Any]] = None) -> Any:
        return self.proxy(session_id, path, method="GET", params=params)

    def post(self, session_id: str, path: str,
             json_data: Optional[Dict[str, Any]] = None,
             data: Optional[Dict[str, Any]] = None) -> Any:
        return self.proxy(session_id, path, method="POST", json_data=json_data, data=data)

    def put(self, session_id: str, path: str,
            json_data: Optional[Dict[str, Any]] = None,
            data: Optional[Dict[str, Any]] = None) -> Any:
        return self.proxy(session_id, path, method="PUT", json_data=json_data, data=data)

    def patch(self, session_id: str, path: str,
              json_data: Optional[Dict[str, Any]] = None,
              data: Optional[Dict[str, Any]] = None) -> Any:
        return self.proxy(session_id, path, method="PATCH", json_data=json_data, data=data)

    def delete(self, session_id: str, path: str,
               params: Optional[Dict[str, Any]] = None) -> Any:
        return self.proxy(session_id, path, method="DELETE", params=params)

    def options(self, session_id: str, path: str,
                params: Optional[Dict[str, Any]] = None) -> Any:
        return self.proxy(session_id, path, method="OPTIONS", params=params)

    def head(self, session_id: str, path: str,
             params: Optional[Dict[str, Any]] = None) -> Any:
        return self.proxy(session_id, path, method="HEAD", params=params)

    def vnc_websockify(self, session_id: str) -> Dict[str, Any]:
        """Get the VNC WebSocket endpoint info for a computer-use session."""
        safe_session_id = _encode_path_segment(session_id)
        return self.client._request_json(
            "GET", f"/v1/computeruse/{safe_session_id}/vnc/websockify"
        )


class APIKeysManager:
    """API key management wrapper."""

    def __init__(self, client: InstaVM):
        self.client = client

    def create(self, **payload: Any) -> Dict[str, Any]:
        return self.client._request_json("POST", "/v1/api-keys/", json_data=payload)

    def list(self) -> List[Dict[str, Any]]:
        result = self.client._request_json("GET", "/v1/api-keys/")
        return result if isinstance(result, list) else []

    def get(self, item_id: int) -> Dict[str, Any]:
        safe_id = _encode_path_segment(item_id)
        return self.client._request_json("GET", f"/v1/api-keys/{safe_id}")

    def update(self, item_id: int, **payload: Any) -> Dict[str, Any]:
        safe_id = _encode_path_segment(item_id)
        return self.client._request_json(
            "PATCH",
            f"/v1/api-keys/{safe_id}",
            json_data=payload
        )

    def delete(self, item_id: int) -> Dict[str, Any]:
        safe_id = _encode_path_segment(item_id)
        return self.client._request_json("DELETE", f"/v1/api-keys/{safe_id}")


class AuditManager:
    """Audit log wrapper."""

    def __init__(self, client: InstaVM):
        self.client = client

    def catalog(self) -> Dict[str, Any]:
        return self.client._request_json("GET", "/v1/audit/catalog")

    def events(self, event_group: Optional[str] = None, event_type: Optional[str] = None,
               status: Optional[str] = None, user_id: Optional[int] = None,
               vm_id: Optional[str] = None, session_id: Optional[str] = None,
               cursor: Optional[str] = None, limit: Optional[int] = None) -> Dict[str, Any]:
        params: Dict[str, Any] = {}
        if event_group is not None:
            params["event_group"] = event_group
        if event_type is not None:
            params["event_type"] = event_type
        if status is not None:
            params["status"] = status
        if user_id is not None:
            params["user_id"] = user_id
        if vm_id is not None:
            params["vm_id"] = vm_id
        if session_id is not None:
            params["session_id"] = session_id
        if cursor is not None:
            params["cursor"] = cursor
        if limit is not None:
            params["limit"] = limit
        return self.client._request_json("GET", "/v1/audit/events", params=params)

    def get_event(self, event_id: str) -> Dict[str, Any]:
        safe_event_id = _encode_path_segment(event_id)
        return self.client._request_json("GET", f"/v1/audit/events/{safe_event_id}")


class WebhooksManager:
    """Webhook endpoint + delivery wrapper."""

    def __init__(self, client: InstaVM):
        self.client = client

    def create_endpoint(self, url: str, description: Optional[str] = None,
                        event_patterns: Optional[List[str]] = None,
                        timeout_seconds: Optional[int] = None,
                        max_retries: Optional[int] = None) -> Dict[str, Any]:
        body: Dict[str, Any] = {"url": url}
        if description is not None:
            body["description"] = description
        if event_patterns is not None:
            body["event_patterns"] = event_patterns
        if timeout_seconds is not None:
            body["timeout_seconds"] = timeout_seconds
        if max_retries is not None:
            body["max_retries"] = max_retries
        return self.client._request_json("POST", "/v1/webhooks/endpoints", json_data=body)

    def list_endpoints(self) -> List[Dict[str, Any]]:
        result = self.client._request_json("GET", "/v1/webhooks/endpoints")
        return result if isinstance(result, list) else []

    def get_endpoint(self, endpoint_id: str) -> Dict[str, Any]:
        safe_endpoint_id = _encode_path_segment(endpoint_id)
        return self.client._request_json("GET", f"/v1/webhooks/endpoints/{safe_endpoint_id}")

    def update_endpoint(self, endpoint_id: str, **payload: Any) -> Dict[str, Any]:
        safe_endpoint_id = _encode_path_segment(endpoint_id)
        return self.client._request_json(
            "PATCH",
            f"/v1/webhooks/endpoints/{safe_endpoint_id}",
            json_data=payload
        )

    def delete_endpoint(self, endpoint_id: str) -> Dict[str, Any]:
        safe_endpoint_id = _encode_path_segment(endpoint_id)
        return self.client._request_json("DELETE", f"/v1/webhooks/endpoints/{safe_endpoint_id}")

    def rotate_secret(self, endpoint_id: str) -> Dict[str, Any]:
        safe_endpoint_id = _encode_path_segment(endpoint_id)
        return self.client._request_json("POST", f"/v1/webhooks/endpoints/{safe_endpoint_id}/rotate-secret")

    def send_test_event(self, endpoint_id: str) -> Dict[str, Any]:
        safe_endpoint_id = _encode_path_segment(endpoint_id)
        return self.client._request_json("POST", f"/v1/webhooks/endpoints/{safe_endpoint_id}/test")

    def verify_endpoint(self, endpoint_id: str) -> Dict[str, Any]:
        safe_endpoint_id = _encode_path_segment(endpoint_id)
        return self.client._request_json("POST", f"/v1/webhooks/endpoints/{safe_endpoint_id}/verify")

    def list_deliveries(self, endpoint_id: Optional[str] = None, status: Optional[str] = None,
                        event_type: Optional[str] = None, cursor: Optional[str] = None,
                        limit: Optional[int] = None) -> Dict[str, Any]:
        params: Dict[str, Any] = {}
        if endpoint_id is not None:
            params["endpoint_id"] = endpoint_id
        if status is not None:
            params["status"] = status
        if event_type is not None:
            params["event_type"] = event_type
        if cursor is not None:
            params["cursor"] = cursor
        if limit is not None:
            params["limit"] = limit
        return self.client._request_json("GET", "/v1/webhooks/deliveries", params=params)

    def replay_delivery(self, delivery_id: str) -> Dict[str, Any]:
        safe_delivery_id = _encode_path_segment(delivery_id)
        return self.client._request_json("POST", f"/v1/webhooks/deliveries/{safe_delivery_id}/replay")
