# Copyright 2025 Daytona Platforms Inc.
# SPDX-License-Identifier: Apache-2.0

# DO NOT EDIT THIS FILE MANUALLY.
# This file is auto-generated by the unasync conversion script.
# Edit the async source and re-run this script.

import asyncio
import json
import time
import warnings
from typing import Callable, Dict, List, Optional, Union, overload

try:
    from importlib.metadata import version
except ImportError:
    # Fallback for Python < 3.8
    from importlib_metadata import version

from daytona_api_client import (
    ApiClient,
    Configuration,
    CreateBuildInfo,
    CreateSandbox,
    ObjectStorageApi,
    SandboxApi,
    SandboxState,
    SnapshotsApi,
)
from daytona_api_client import ToolboxApi as ToolboxApi
from daytona_api_client import VolumesApi as VolumesApi
from environs import Env

from .._utils.enum import to_enum
from .._utils.errors import DaytonaError, intercept_errors
from .._utils.stream import process_streaming_response
from .._utils.timeout import with_timeout
from ..code_toolbox.sandbox_python_code_toolbox import SandboxPythonCodeToolbox
from ..code_toolbox.sandbox_ts_code_toolbox import SandboxTsCodeToolbox
from ..common.daytona import (
    CodeLanguage,
    CreateSandboxFromImageParams,
    CreateSandboxFromSnapshotParams,
    DaytonaConfig,
    Image,
)
from .sandbox import Sandbox
from .snapshot import SnapshotService
from .volume import VolumeService


class Daytona:
    """Main class for interacting with the Daytona API.

    This class provides methods to create, manage, and interact with Daytona Sandboxes.
    It can be initialized either with explicit configuration or using environment variables.

    Attributes:
        volume (VolumeService): Service for managing volumes.
        snapshot (SnapshotService): Service for managing snapshots.

    Example:
        Using environment variables:
        ```python
        daytona = Daytona()  # Uses DAYTONA_API_KEY, DAYTONA_API_URL
        sandbox = daytona.create()
        ```

        Using explicit configuration:
        ```python
        config = DaytonaConfig(
            api_key="your-api-key",
            api_url="https://your-api.com",
            target="us"
        )
        daytona = Daytona(config)
        sandbox = daytona.create()
        ```
    """

    def __init__(self, config: Optional[DaytonaConfig] = None):
        """Initializes Daytona instance with optional configuration.

        If no config is provided, reads from environment variables:
        - `DAYTONA_API_KEY`: Required API key for authentication
        - `DAYTONA_API_URL`: Required api URL
        - `DAYTONA_TARGET`: Optional target environment (defaults to 'us')

        Args:
            config (Optional[DaytonaConfig]): Object containing api_key, api_url, and target.

        Raises:
            DaytonaError: If API key is not provided either through config or environment variables

        Example:
            ```python
            from daytona import Daytona, DaytonaConfig
            # Using environment variables
            daytona1 = Daytona()

            # Using explicit configuration
            config = DaytonaConfig(
                api_key="your-api-key",
                api_url="https://your-api.com",
                target="us"
            )
            daytona2 = Daytona(config)

            ```
        """

        default_api_url = "https://app.daytona.io/api"
        self.default_language = CodeLanguage.PYTHON

        if config is None or (
            not all([config.api_key, config.api_url, config.target])
            and not all(
                [
                    config.jwt_token,
                    config.organization_id,
                    config.api_url,
                    config.target,
                ]
            )
        ):
            # Initialize env - it automatically reads from .env and .env.local
            env = Env()
            env.read_env()  # reads .env
            # reads .env.local and overrides values
            env.read_env(".env.local", override=True)

            self._api_key = env.str("DAYTONA_API_KEY", None)
            self._jwt_token = env.str("DAYTONA_JWT_TOKEN", None)
            self._organization_id = env.str("DAYTONA_ORGANIZATION_ID", None)
            self._api_url = env.str("DAYTONA_API_URL", None) or env.str("DAYTONA_SERVER_URL", default_api_url)
            self._target = env.str("DAYTONA_TARGET", None)

            if env.str("DAYTONA_SERVER_URL", None) and not env.str("DAYTONA_API_URL", None):
                warnings.warn(
                    "Environment variable `DAYTONA_SERVER_URL` is deprecated and will be removed in future versions. "
                    + "Use `DAYTONA_API_URL` instead.",
                    DeprecationWarning,
                    stacklevel=2,
                )

        if config:
            if not config.api_key and config.jwt_token:
                self._api_key = None
            else:
                self._api_key = config.api_key or getattr(self, "api_key", None)
            self._jwt_token = config.jwt_token or getattr(self, "jwt_token", None)
            self._organization_id = config.organization_id or getattr(self, "organization_id", None)
            self._api_url = config.api_url or self._api_url
            self._target = config.target or self._target

        if not self._api_key and not self._jwt_token:
            raise DaytonaError("API key or JWT token is required")

        # Create API configuration without api_key
        configuration = Configuration(host=self._api_url)
        self._api_client = ApiClient(configuration)
        self._api_client.default_headers["Authorization"] = f"Bearer {self._api_key or self._jwt_token}"
        self._api_client.default_headers["X-Daytona-Source"] = "python-sdk"

        # Get SDK version dynamically
        try:
            sdk_version = None
            for pkg_name in ["daytona_sdk", "daytona"]:
                try:
                    sdk_version = version(pkg_name)
                    break
                except Exception:
                    continue

            if sdk_version is None:
                raise Exception("Neither package found")

        except Exception:
            # Fallback version if neither package metadata is available
            sdk_version = "unknown"
        self._api_client.default_headers["X-Daytona-SDK-Version"] = sdk_version

        if not self._api_key:
            if not self._organization_id:
                raise DaytonaError("Organization ID is required when using JWT token")
            self._api_client.default_headers["X-Daytona-Organization-ID"] = self._organization_id

        # Initialize API clients with the api_client instance
        self._sandbox_api = SandboxApi(self._api_client)
        self._toolbox_api = ToolboxApi(self._api_client)
        self._object_storage_api = ObjectStorageApi(self._api_client)

        # Initialize services
        self.volume = VolumeService(VolumesApi(self._api_client))
        self.snapshot = SnapshotService(SnapshotsApi(self._api_client), self._object_storage_api)

    @overload
    def create(
        self,
        params: Optional[CreateSandboxFromSnapshotParams] = None,
        *,
        timeout: Optional[float] = 60,
    ) -> Sandbox:
        """Creates Sandboxes from specified or default snapshot. You can specify various parameters,
        including language, image, environment variables, and volumes.

        Args:
            params (Optional[CreateSandboxFromSnapshotParams]): Parameters for Sandbox creation. If not provided,
                   defaults to default Daytona snapshot and Python language.
            timeout (Optional[float]): Timeout (in seconds) for sandbox creation. 0 means no timeout.
                Default is 60 seconds.

        Returns:
            Sandbox: The created Sandbox instance.

        Raises:
            DaytonaError: If timeout, auto_stop_interval or auto_archive_interval is negative;
                If sandbox fails to start or times out

        Example:
            Create a default Python Sandbox:
            ```python
            sandbox = daytona.create()
            ```

            Create a custom Sandbox:
            ```python
            params = CreateSandboxFromSnapshotParams(
                language="python",
                snapshot="my-snapshot-id",
                env_vars={"DEBUG": "true"},
                auto_stop_interval=0,
                auto_archive_interval=60
            )
            sandbox = daytona.create(params, timeout=40)
            ```
        """

    @overload
    def create(
        self,
        params: Optional[CreateSandboxFromImageParams] = None,
        *,
        timeout: Optional[float] = 60,
        on_snapshot_create_logs: Callable[[str], None] = None,
    ) -> Sandbox:
        """Creates Sandboxes from specified image available on some registry or declarative Daytona Image.
        You can specify various parameters, including resources, language, image, environment variables,
        and volumes. Daytona creates snapshot from provided image and uses it to create Sandbox.

        Args:
            params (Optional[CreateSandboxFromImageParams]): Parameters for Sandbox creation from image.
            timeout (Optional[float]): Timeout (in seconds) for sandbox creation. 0 means no timeout.
                Default is 60 seconds.
            on_snapshot_create_logs (Callable[[str], None]): This callback function handles snapshot creation logs.

        Returns:
            Sandbox: The created Sandbox instance.

        Raises:
            DaytonaError: If timeout, auto_stop_interval or auto_archive_interval is negative;
                If sandbox fails to start or times out

        Example:
            Create a default Python Sandbox from image:
            ```python
            sandbox = daytona.create(CreateSandboxFromImageParams(image="debian:12.9"))
            ```

            Create a custom Sandbox from declarative Image definition:
            ```python
            declarative_image = (
                Image.base("alpine:3.18")
                .pipInstall(["numpy", "pandas"])
                .env({"MY_ENV_VAR": "My Environment Variable"})
            )
            params = CreateSandboxFromImageParams(
                language="python",
                image=declarative_image,
                env_vars={"DEBUG": "true"},
                resources=Resources(cpu=2, memory=4),
                auto_stop_interval=0,
                auto_archive_interval=60,
            )
            sandbox = daytona.create(
                params,
                timeout=40,
                on_snapshot_create_logs=lambda chunk: print(chunk, end=""),
            )
            ```
        """

    @intercept_errors(message_prefix="Failed to create sandbox: ")
    def create(
        self,
        params: Optional[Union[CreateSandboxFromSnapshotParams, CreateSandboxFromImageParams]] = None,
        *,
        timeout: Optional[float] = 60,
        on_snapshot_create_logs: Callable[[str], None] = None,
    ) -> Sandbox:
        # If no params provided, create default params for Python
        if params is None:
            params = CreateSandboxFromSnapshotParams(language=self.default_language)
        elif params.language is None:
            params.language = self.default_language

        return self._create(params, timeout=timeout, on_snapshot_create_logs=on_snapshot_create_logs)

    @with_timeout(
        error_message=lambda self, timeout: (
            f"Failed to create and start sandbox within {timeout} seconds timeout period."
        )
    )
    def _create(
        self,
        params: Optional[Union[CreateSandboxFromSnapshotParams, CreateSandboxFromImageParams]] = None,
        *,
        timeout: Optional[float] = 60,
        on_snapshot_create_logs: Callable[[str], None] = None,
    ) -> Sandbox:
        code_toolbox = self._get_code_toolbox(params.language)

        if timeout < 0:
            raise DaytonaError("Timeout must be a non-negative number")

        if params.auto_stop_interval is not None and params.auto_stop_interval < 0:
            raise DaytonaError("auto_stop_interval must be a non-negative integer")

        if params.auto_archive_interval is not None and params.auto_archive_interval < 0:
            raise DaytonaError("auto_archive_interval must be a non-negative integer")

        target = self._target

        # Create sandbox using dictionary
        sandbox_data = CreateSandbox(
            user=params.os_user,
            env=params.env_vars if params.env_vars else {},
            labels=params.labels,
            public=params.public,
            target=str(target) if target else None,
            auto_stop_interval=params.auto_stop_interval,
            auto_archive_interval=params.auto_archive_interval,
            volumes=params.volumes,
        )

        if getattr(params, "snapshot", None):
            sandbox_data.snapshot = params.snapshot

        if getattr(params, "image", None):
            if isinstance(params.image, str):
                sandbox_data.build_info = CreateBuildInfo(
                    dockerfile_content=Image.base(params.image).dockerfile(),
                )
            elif isinstance(params.image, Image):
                context_hashes = SnapshotService.process_image_context(self._object_storage_api, params.image)
                sandbox_data.build_info = CreateBuildInfo(
                    context_hashes=context_hashes,
                    dockerfile_content=params.image.dockerfile(),
                )

        if getattr(params, "resources", None):
            sandbox_data.cpu = params.resources.cpu
            sandbox_data.memory = params.resources.memory
            sandbox_data.disk = params.resources.disk
            sandbox_data.gpu = params.resources.gpu

        response = self._sandbox_api.create_sandbox(sandbox_data, _request_timeout=timeout or None)

        if response.state == SandboxState.PENDING_BUILD and on_snapshot_create_logs:
            _, url, *_ = self._sandbox_api._get_build_logs_serialize(  # pylint: disable=protected-access
                response.id,
                follow=True,
                x_daytona_organization_id=None,
                _request_auth=None,
                _content_type=None,
                _headers=None,
                _host_index=None,
            )

            response_ref = {"response": response}

            def should_terminate():
                response_ref["response"] = self._sandbox_api.get_sandbox(response_ref["response"].id)
                return response_ref["response"].state in [
                    SandboxState.STARTED,
                    SandboxState.STARTING,
                    SandboxState.ERROR,
                    SandboxState.BUILD_FAILED,
                ]

            while response_ref["response"].state == SandboxState.PENDING_BUILD:
                time.sleep(1)
                response_ref["response"] = self._sandbox_api.get_sandbox(response_ref["response"].id)

            asyncio.run(
                process_streaming_response(
                    url=url,
                    headers=self._sandbox_api.api_client.default_headers,
                    on_chunk=lambda chunk: on_snapshot_create_logs(chunk.rstrip()),
                    should_terminate=should_terminate,
                )
            )
            response = response_ref["response"]

        sandbox = Sandbox(
            response,
            self._sandbox_api,
            self._toolbox_api,
            code_toolbox,
        )

        if sandbox.state != SandboxState.STARTED:
            # Wait for sandbox to start
            try:
                sandbox.wait_for_sandbox_start()
            finally:
                # If not Daytona SaaS, we don't need to handle pulling image state
                pass

        return sandbox

    def _get_code_toolbox(self, language: Optional[CodeLanguage] = None):
        """Helper method to get the appropriate code toolbox based on language.

        Args:
            language (Optional[CodeLanguage]): Language of the code toolbox. If not provided, defaults to Python.

        Returns:
            The appropriate code toolbox instance for the specified language.

        Raises:
            DaytonaError: If an unsupported language is specified.
        """
        if not language:
            return SandboxPythonCodeToolbox()

        enum_language = to_enum(CodeLanguage, language)
        if enum_language is None:
            raise DaytonaError(f"Unsupported language: {language}")
        language = enum_language

        match language:
            case CodeLanguage.JAVASCRIPT | CodeLanguage.TYPESCRIPT:
                return SandboxTsCodeToolbox()
            case CodeLanguage.PYTHON:
                return SandboxPythonCodeToolbox()
            case _:
                raise DaytonaError(f"Unsupported language: {language}")

    def delete(self, sandbox: Sandbox, timeout: Optional[float] = 60) -> None:
        """Deletes a Sandbox.

        Args:
            sandbox (Sandbox): The Sandbox instance to delete.
            timeout (Optional[float]): Timeout (in seconds) for sandbox deletion. 0 means no timeout.
                Default is 60 seconds.

        Raises:
            DaytonaError: If sandbox fails to delete or times out

        Example:
            ```python
            sandbox = daytona.create()
            # ... use sandbox ...
            daytona.delete(sandbox)  # Clean up when done
            ```
        """
        return sandbox.delete(timeout)

    @intercept_errors(message_prefix="Failed to get sandbox: ")
    def get(self, sandbox_id: str) -> Sandbox:
        """Gets a Sandbox by its ID.

        Args:
            sandbox_id (str): The ID of the Sandbox to retrieve.

        Returns:
            Sandbox: The Sandbox instance.

        Raises:
            DaytonaError: If sandbox_id is not provided.

        Example:
            ```python
            sandbox = daytona.get("my-sandbox-id")
            print(sandbox.status)
            ```
        """
        if not sandbox_id:
            raise DaytonaError("sandbox_id is required")

        # Get the sandbox instance
        sandbox_instance = self._sandbox_api.get_sandbox(sandbox_id)

        # Create and return sandbox with Python code toolbox as default
        code_toolbox = SandboxPythonCodeToolbox()
        return Sandbox(
            sandbox_instance,
            self._sandbox_api,
            self._toolbox_api,
            code_toolbox,
        )

    @intercept_errors(message_prefix="Failed to find sandbox: ")
    def find_one(self, sandbox_id: Optional[str] = None, labels: Optional[Dict[str, str]] = None) -> Sandbox:
        """Finds a Sandbox by its ID or labels.

        Args:
            sandbox_id (Optional[str]): The ID of the Sandbox to retrieve.
            labels (Optional[Dict[str, str]]): Labels to filter Sandboxes.

        Returns:
            Sandbox: First Sandbox that matches the ID or labels.

        Raises:
            DaytonaError: If no Sandbox is found.

        Example:
            ```python
            sandbox = daytona.find_one(labels={"my-label": "my-value"})
            print(f"Sandbox ID: {sandbox.id} State: {sandbox.state}")
            ```
        """
        if sandbox_id:
            return self.get(sandbox_id)
        sandboxes = self.list(labels)
        if len(sandboxes) == 0:
            raise DaytonaError(f"No sandbox found with labels {labels}")
        return sandboxes[0]

    @intercept_errors(message_prefix="Failed to list sandboxes: ")
    def list(self, labels: Optional[Dict[str, str]] = None) -> List[Sandbox]:
        """Lists Sandboxes filtered by labels.

        Args:
            labels (Optional[Dict[str, str]]): Labels to filter Sandboxes.

        Returns:
            List[Sandbox]: List of Sandbox instances that match the labels.

        Example:
            ```python
            sandboxes = daytona.list(labels={"my-label": "my-value"})
            for sandbox in sandboxes:
                print(f"{sandbox.id}: {sandbox.status}")
            ```
        """
        sandboxes = self._sandbox_api.list_sandboxes(labels=json.dumps(labels))

        return [
            Sandbox(
                sandbox,
                self._sandbox_api,
                self._toolbox_api,
                self._get_code_toolbox(self._validate_language_label(sandbox.labels.get("code-toolbox-language"))),
            )
            for sandbox in sandboxes
        ]

    def _validate_language_label(self, language: Optional[str]) -> CodeLanguage:
        """Validates and normalizes the language label.

        Args:
            language (Optional[str]): The language label to validate.

        Returns:
            CodeLanguage: The validated language, defaults to "python" if None

        Raises:
            DaytonaError: If the language is not supported.
        """
        if not language:
            return CodeLanguage.PYTHON

        enum_language = to_enum(CodeLanguage, language)
        if enum_language is None:
            raise DaytonaError(f"Invalid code-toolbox-language: {language}")
        return enum_language

    def start(self, sandbox: Sandbox, timeout: Optional[float] = 60) -> None:
        """Starts a Sandbox and waits for it to be ready.

        Args:
            sandbox (Sandbox): The Sandbox to start.
            timeout (Optional[float]): Optional timeout in seconds to wait for the Sandbox to start.
                0 means no timeout. Default is 60 seconds.

        Raises:
            DaytonaError: If timeout is negative; If Sandbox fails to start or times out
        """
        sandbox.start(timeout)

    def stop(self, sandbox: Sandbox, timeout: Optional[float] = 60) -> None:
        """Stops a Sandbox and waits for it to be stopped.

        Args:
            sandbox (Sandbox): The sandbox to stop
            timeout (Optional[float]): Optional timeout (in seconds) for sandbox stop.
                0 means no timeout. Default is 60 seconds.

        Raises:
            DaytonaError: If timeout is negative; If Sandbox fails to stop or times out
        """
        sandbox.stop(timeout)
