import asyncio
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


class DaytonaEnvironment(BaseEnvironment):
    def __init__(
        self,
        environment_dir: Path,
        environment_name: str,
        session_id: str,
        trial_paths: TrialPaths,
        task_env_config: EnvironmentConfig,
        *args,
        **kwargs,
    ):
        super().__init__(
            environment_dir=environment_dir,
            environment_name=environment_name,
            session_id=session_id,
            trial_paths=trial_paths,
            task_env_config=task_env_config,
            *args,
            **kwargs,
        )

        self._workdir = next(
            (
                line.strip().split(maxsplit=1)[1]
                for line in reversed(
                    self._environment_definition_path.read_text().splitlines()
                )
                if line.strip().upper().startswith("WORKDIR")
                and len(line.strip().split(maxsplit=1)) == 2
            ),
            None,
        )  # TODO(alexgshaw) this should be unnecessary by Sept 2025

        self._sandbox: AsyncSandbox | None = None
        self._daytona: AsyncDaytona | None = None

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

    @property
    def is_mounted(self) -> bool:
        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._daytona:
            raise RuntimeError("Daytona not found. This should never happen.")

        self._sandbox = await self._daytona.create(
            params=params, timeout=round(self.task_env_config.build_timeout_sec)
        )

    async def start(self, force_build: bool):
        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._daytona = AsyncDaytona()

        if not force_build:
            try:
                snapshot = await self._daytona.snapshot.get(
                    self.environment_dir.parent.name
                )
                if snapshot.state != SnapshotState.ACTIVE:
                    force_build = True
            except Exception:
                force_build = True

        if force_build:
            image = Image.from_dockerfile(self._environment_definition_path)

            params = CreateSandboxFromImageParams(
                image=image,
                auto_delete_interval=0,
                resources=resources,
            )
        else:
            params = CreateSandboxFromSnapshotParams(
                snapshot=self.environment_dir.parent.name,
            )

        await self._create_sandbox(params=params)

        await self.exec(f"mkdir -p {EnvironmentPaths.agent_dir}")
        await self.exec(f"mkdir -p {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."
            )

        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

        if self._daytona:
            try:
                await self._daytona.close()
            except Exception as e:
                self.logger.error(f"Error closing daytona: {e}")
            finally:
                self._daytona = 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}"

            response = await self._sandbox.process.execute_session_command(
                session_id,
                SessionExecuteRequest(
                    command=f"cd {cwd or self._workdir} && {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)
                pass
            except Exception as e:
                self.logger.warning(f"Could not delete session: {e}")

        return result
