import asyncio
import atexit
import os
import shlex
from pathlib import Path
from uuid import uuid4

from daytona import (
    AsyncDaytona,
    AsyncSandbox,
    CreateSandboxFromImageParams,
    CreateSandboxFromSnapshotParams,
    FileDownloadRequest,
    FileUpload,
    Image,
    Resources,
    SessionExecuteRequest,
)
from daytona._async.snapshot import SnapshotState
from tenacity import retry, stop_after_attempt, wait_exponential

from harbor.environments.base import BaseEnvironment, ExecResult
from harbor.models.environment_type import EnvironmentType
from harbor.models.task.config import EnvironmentConfig
from harbor.models.trial.paths import EnvironmentPaths, TrialPaths
from harbor.utils.logger import logger


class DaytonaClientManager:
    """
    Singleton manager for the AsyncDaytona client.

    Ensures a single shared client instance across all DaytonaEnvironment instances,
    with proper cleanup at program termination.
    """

    _instance: "DaytonaClientManager | None" = None
    _lock = asyncio.Lock()

    def __init__(self):
        self._client: AsyncDaytona | None = None
        self._client_lock = asyncio.Lock()
        self._logger = logger.getChild(__name__)
        self._cleanup_registered = False

    @classmethod
    async def get_instance(cls) -> "DaytonaClientManager":
        """Get or create the singleton instance."""
        if cls._instance is None:
            async with cls._lock:
                if cls._instance is None:
                    cls._instance = cls()

        assert cls._instance is not None

        return cls._instance

    async def get_client(self) -> AsyncDaytona:
        """
        Get the shared AsyncDaytona client, creating it if necessary.
        """
        async with self._client_lock:
            if self._client is None:
                self._logger.debug("Creating new AsyncDaytona client")
                self._client = AsyncDaytona()

                # Register cleanup handler on first client creation
                if not self._cleanup_registered:
                    atexit.register(self._cleanup_sync)
                    self._cleanup_registered = True

            return self._client

    def _cleanup_sync(self):
        """Synchronous cleanup wrapper for atexit."""
        try:
            asyncio.run(self._cleanup())
        except Exception as e:
            # Use print since logging might not be available during shutdown
            print(f"Error during Daytona client cleanup: {e}")

    async def _cleanup(self):
        """Close the Daytona client if it exists."""
        async with self._client_lock:
            if self._client is not None:
                try:
                    self._logger.debug("Closing AsyncDaytona client at program exit")
                    await self._client.close()
                    self._logger.debug("AsyncDaytona client closed successfully")
                except Exception as e:
                    self._logger.error(f"Error closing AsyncDaytona client: {e}")
                finally:
                    self._client = None


class DaytonaEnvironment(BaseEnvironment):
    def __init__(
        self,
        environment_dir: Path,
        environment_name: str,
        session_id: str,
        trial_paths: TrialPaths,
        task_env_config: EnvironmentConfig,
        snapshot_template_name: str | None = None,
        network_block_all: bool = False,
        **kwargs,
    ):
        """
        Initialize a DaytonaEnvironment instance.

        Args:
            environment_dir: The directory containing the environment definition files.
                Must contain a Dockerfile for building the environment.
            environment_name: The name identifier for this environment instance.
            session_id: Unique session identifier for this environment instance.
            trial_paths: Path configuration object containing trial-related directory paths.
            task_env_config: Environment configuration specifying resources (CPU, memory,
                storage), build timeout, and optional Docker image.
            snapshot_template_name: Optional name of a snapshot template to use for
                faster environment initialization. Should contain the {name} placeholder.
                For example, "harbor__{name}__snapshot". The placeholder will be replaced
                with the environment_name.
            network_block_all: Whether to block all network access for the sandbox.
                Useful for air-gapped environments.

        Raises:
            FileNotFoundError: If the Dockerfile is not found in the environment_dir.
        """

        super().__init__(
            environment_dir=environment_dir,
            environment_name=environment_name,
            session_id=session_id,
            trial_paths=trial_paths,
            task_env_config=task_env_config,
            **kwargs,
        )

        self._snapshot_template_name = snapshot_template_name
        self._network_block_all = network_block_all

        self._sandbox: AsyncSandbox | None = None
        self._client_manager: DaytonaClientManager | None = None

    @staticmethod
    def type() -> EnvironmentType:
        return EnvironmentType.DAYTONA

    @property
    def is_mounted(self) -> bool:
        return False

    @staticmethod
    def supports_gpus() -> bool:
        """Daytona does not currently support GPU allocation."""
        return False

    @property
    def _environment_definition_path(self) -> Path:
        return self.environment_dir / "Dockerfile"

    def _validate_definition(self):
        if not self._environment_definition_path.exists():
            raise FileNotFoundError(
                f"{self._environment_definition_path} not found. Please ensure the "
                "file exists."
            )

    @retry(
        stop=stop_after_attempt(2),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def _create_sandbox(
        self, params: CreateSandboxFromImageParams | CreateSandboxFromSnapshotParams
    ):
        if not self._client_manager:
            raise RuntimeError(
                "Client manager not initialized. This should never happen."
            )

        daytona = await self._client_manager.get_client()
        self._sandbox = await daytona.create(
            params=params, timeout=round(self.task_env_config.build_timeout_sec)
        )

    async def start(self, force_build: bool) -> None:
        """
        Start the environment and optionally forces a build. Starts the sandbox using
        the following preference order:

        1. Snapshot if it exists and snapshot template is specified.
        2. Use prebuilt image from task environment config if it exists.
        2. Build from Dockerfile if force_build is True or no docker image in config.

        Args:
            force_build: Whether to force a build.

        Raises:
            RuntimeError: If the environment is not found.
        """

        resources = Resources(
            cpu=self.task_env_config.cpus,
            memory=self.task_env_config.memory_mb // 1024,
            disk=self.task_env_config.storage_mb // 1024,
        )

        self._client_manager = await DaytonaClientManager.get_instance()
        daytona = await self._client_manager.get_client()

        snapshot_name: str | None = None
        snapshot_exists = False

        if self._snapshot_template_name:
            snapshot_name = self._snapshot_template_name.format(
                name=self.environment_name
            )

            try:
                snapshot = await daytona.snapshot.get(snapshot_name)
                if snapshot.state == SnapshotState.ACTIVE:
                    snapshot_exists = True
            except Exception:
                snapshot_exists = False

        if snapshot_exists and force_build:
            self.logger.warning(
                "Snapshot template specified but force_build is True. "
                "Snapshot will be used instead of building from scratch."
            )

        if snapshot_exists and snapshot_name:
            self.logger.debug(f"Using snapshot: {snapshot_name}")
            params = CreateSandboxFromSnapshotParams(
                auto_delete_interval=0,
                snapshot=snapshot_name,
                network_block_all=self._network_block_all,
            )
        elif force_build or not self.task_env_config.docker_image:
            self.logger.debug(
                f"Building environment from Dockerfile {
                    self._environment_definition_path
                }"
            )
            image = Image.from_dockerfile(self._environment_definition_path)
            params = CreateSandboxFromImageParams(
                image=image,
                auto_delete_interval=0,
                resources=resources,
                network_block_all=self._network_block_all,
            )
        else:
            self.logger.debug(
                f"Using prebuilt image: {self.task_env_config.docker_image}"
            )
            image = Image.base(self.task_env_config.docker_image)
            params = CreateSandboxFromImageParams(
                image=image,
                auto_delete_interval=0,
                resources=resources,
                network_block_all=self._network_block_all,
            )

        await self._create_sandbox(params=params)

        await self.exec(
            f"mkdir -p {str(EnvironmentPaths.agent_dir)} {
                str(EnvironmentPaths.verifier_dir)
            }"
        )

    @retry(
        stop=stop_after_attempt(2),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def _stop_sandbox(self):
        if self._sandbox:
            await self._sandbox.delete()

    async def stop(self, delete: bool):
        if not delete:
            self.logger.info(
                "Daytona harbor are ephemeral and will be deleted after use, "
                "regardless of delete=False."
            )

        try:
            if not self._sandbox:
                self.logger.warning(
                    "Sandbox not found. Please build the environment first."
                )
            else:
                try:
                    await self._stop_sandbox()
                except Exception as e:
                    self.logger.error(f"Error stopping sandbox {self._sandbox.id}: {e}")
                finally:
                    self._sandbox = None
        finally:
            # Client cleanup happens at program exit via atexit
            self._client_manager = None

    @retry(
        stop=stop_after_attempt(2),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def upload_file(self, source_path: Path | str, target_path: str):
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please build the environment first.")

        await self._sandbox.fs.upload_file(str(source_path), target_path)

    @retry(
        stop=stop_after_attempt(2),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def upload_dir(self, source_dir: Path | str, target_dir: str):
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please build the environment first.")

        file_uploads = []
        source_dir = Path(source_dir)

        for file_path in source_dir.rglob("*"):
            if file_path.is_file():
                relative_path = file_path.relative_to(Path(source_dir))
                destination_path = str(Path(target_dir) / relative_path)

                file_uploads.append(
                    FileUpload(
                        source=str(file_path),
                        destination=destination_path,
                    )
                )

        if file_uploads:
            await self._sandbox.fs.upload_files(files=file_uploads)

    @retry(
        stop=stop_after_attempt(2),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def download_file(self, source_path: str, target_path: Path | str):
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please build the environment first.")

        await self._sandbox.fs.download_file(source_path, str(target_path))

    @retry(
        stop=stop_after_attempt(2),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def download_dir(self, source_dir: str, target_dir: Path | str):
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please build the environment first.")

        target_dir = Path(target_dir)
        target_dir.mkdir(parents=True, exist_ok=True)

        search_result = await self._sandbox.fs.search_files(source_dir, "*")

        file_downloads = []
        for file_path in search_result.files:
            file_info = await self._sandbox.fs.get_file_info(file_path)

            if not file_info.is_dir:
                path_obj = Path(file_path)
                relative_path = path_obj.relative_to(Path(source_dir))
                local_file_path = target_dir / relative_path

                local_file_path.parent.mkdir(parents=True, exist_ok=True)

                file_downloads.append(
                    FileDownloadRequest(
                        source=file_path,
                        destination=str(local_file_path),
                    )
                )

        if file_downloads:
            await self._sandbox.fs.download_files(files=file_downloads)

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def _get_session_command_with_retry(self, session_id: str, command_id: str):
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please build the environment first.")
        return await self._sandbox.process.get_session_command(session_id, command_id)

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def _get_session_command_logs_with_retry(
        self, session_id: str, command_id: str
    ):
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please build the environment first.")
        return await self._sandbox.process.get_session_command_logs(
            session_id, command_id
        )

    async def _poll_response(self, session_id: str, command_id: str):
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please build the environment first.")

        response = await self._get_session_command_with_retry(session_id, command_id)

        while response.exit_code is None:
            await asyncio.sleep(1)
            response = await self._get_session_command_with_retry(
                session_id,
                response.id,
            )

        logs = await self._get_session_command_logs_with_retry(session_id, command_id)

        return ExecResult(
            stdout=logs.stdout,
            stderr=logs.stderr,
            return_code=int(response.exit_code),
        )

    async def exec(
        self,
        command: str,
        cwd: str | None = None,
        env: dict[str, str] | None = None,
        timeout_sec: int | None = None,
    ) -> ExecResult:
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please build the environment first.")

        session_id = str(uuid4())
        try:
            await self._sandbox.process.create_session(session_id)

            command = f"bash -ic {shlex.quote(command)}"

            if env:
                for key, value in env.items():
                    command = f"{key}={shlex.quote(value)} {command}"

            if timeout_sec:
                command = f"timeout {timeout_sec} {command}"

            if cwd:
                command = f"cd {cwd} && {command}"

            response = await self._sandbox.process.execute_session_command(
                session_id,
                SessionExecuteRequest(
                    command=command,
                    run_async=True,  # type: ignore
                ),
                timeout=timeout_sec,
            )

            if response.cmd_id is None:
                raise RuntimeError("Cannot find command ID.")

            result = await self._poll_response(session_id, response.cmd_id)

        finally:
            try:
                # await self._sandbox.process.delete_session(session_id) Don't delete
                # the session bc Daytona will delete any child processes, even if they
                # were started in the background.
                pass
            except Exception as e:
                self.logger.warning(f"Could not delete session: {e}")

        return result

    async def attach(self) -> None:
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please start the environment first.")

        ssh_access = await self._sandbox.create_ssh_access()

        os.execvp(
            "ssh",
            [
                "ssh",
                f"{ssh_access.token}@ssh.app.daytona.io",
            ],
        )
