import shlex
from datetime import timedelta
from pathlib import Path
from typing import Any

from runloop_api_client import AsyncRunloopSDK
from runloop_api_client.lib.polling import PollingConfig
from runloop_api_client.types.blueprint_create_params import BuildContext
from runloop_api_client.types.shared_params.launch_parameters import (
    LaunchParameters,
    UserParameters,
)
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 RunloopEnvironment(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,
            **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
            ),
            "/workspace",
        )

        self._devbox: Any | None = None
        self._client: AsyncRunloopSDK | None = None
        self._shell_name: str = "main_shell"

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

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

    @staticmethod
    def supports_gpus() -> 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."
            )

    def _build_launch_parameters(self) -> LaunchParameters:
        """
        Build launch parameters from EnvironmentConfig, converting to Runloop's format.

        For detailed information on resource sizes and other options, see:
        https://docs.runloop.ai/docs/devboxes/configuration/sizes#custom-resource-sizes
        """
        launch_parameters: LaunchParameters = LaunchParameters(
            architecture="x86_64",
            user_parameters=UserParameters(
                username="root",
                uid=0,
            ),
            resource_size_request="CUSTOM_SIZE",
            custom_cpu_cores=self.task_env_config.cpus,
            custom_gb_memory=self.task_env_config.memory_mb // 1024,
            custom_disk_size=self.task_env_config.storage_mb // 1024,
        )

        return launch_parameters

    async def _get_latest_successful_blueprint_id(
        self,
        blueprint_name: str,
    ) -> str | None:
        """
        Return the most recent successfully built blueprint ID for the given name.

        This uses the underlying AsyncRunloop client (`self._client.api`) to list
        blueprints filtered by name, filters to those in `build_complete` status,
        and then prefers the one with the latest create time.
        """
        if not self._client:
            raise RuntimeError("RunLoop client not found. This should never happen.")

        try:
            page = await self._client.api.blueprints.list(name=blueprint_name)
        except Exception as e:
            self.logger.warning("Failed to list blueprints: %s", e, exc_info=True)
            return None

        blueprints = page.blueprints
        if not blueprints:
            return None

        candidates = [
            bp
            for bp in blueprints
            if bp.name == blueprint_name and bp.status == "build_complete"
        ]
        if not candidates:
            return None

        candidates.sort(key=lambda bp: bp.create_time_ms, reverse=True)
        return candidates[0].id

    @retry(
        stop=stop_after_attempt(2),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def _create_devbox(self, force_build: bool):
        if not self._client:
            raise RuntimeError("RunLoop client not found. This should never happen.")

        blueprint_name = f"harbor_{self.environment_name}_blueprint"
        launch_parameters = self._build_launch_parameters()

        blueprint_id: str | None = None

        if not force_build:
            best_blueprint_id = await self._get_latest_successful_blueprint_id(
                blueprint_name=blueprint_name,
            )
            if best_blueprint_id:
                self.logger.info(
                    "Reusing existing Runloop blueprint %s (id=%s) for environment %s",
                    blueprint_name,
                    best_blueprint_id,
                    self.environment_name,
                )
                blueprint_id = best_blueprint_id

        if not blueprint_id:
            # Either force_build is True or no suitable existing blueprint was found: build a new one.
            self.logger.info(
                "Building new Runloop blueprint %s for environment %s (force_build=%s)",
                blueprint_name,
                self.environment_name,
                force_build,
            )

            # Create devbox from Dockerfile by first building a blueprint.
            dockerfile_content = self._environment_definition_path.read_text()

            # Upload the environment directory as a build context.
            storage_object = await self._client.storage_object.upload_from_dir(
                dir_path=self.environment_dir.resolve(),
                name=f"{self.environment_name}_context.tar.gz",
                ttl=timedelta(hours=1),
            )

            build_context = BuildContext(
                object_id=storage_object.id,
                type="object",
            )

            # Allow long-running blueprint builds (e.g., heavy toolchains, QEMU images).
            # The default PollingConfig(max_attempts=120, interval_seconds=1.0) was too
            # short for several environments and caused PollingTimeout errors even though
            # builds were still progressing. Here we extend both the maximum attempts and
            # add an explicit overall timeout to give blueprints more time to finish.
            polling_config = PollingConfig(
                interval_seconds=2.0,
                max_attempts=900,  # up to ~30 minutes with 2s interval
                timeout_seconds=60 * 60,  # hard cap at 1 hour
            )

            # Create a blueprint with the Dockerfile and wait for it to build
            blueprint = await self._client.blueprint.create(
                dockerfile=dockerfile_content,
                name=blueprint_name,
                build_context=build_context,
                launch_parameters=launch_parameters,
                polling_config=polling_config,
            )

            blueprint_id = blueprint.id

        # Create devbox from the selected or newly created blueprint
        self._devbox = await self._client.devbox.create_from_blueprint_id(
            blueprint_id=blueprint_id,
            name=self.environment_name,
            launch_parameters=launch_parameters,
        )

    async def start(self, force_build: bool):
        if not self._client:
            self._client = AsyncRunloopSDK()

        await self._create_devbox(force_build=force_build)

        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 _shutdown_devbox(self):
        if self._devbox and self._client:
            await self._devbox.shutdown()

    async def stop(self, delete: bool):
        if not self._devbox:
            self.logger.warning("Devbox not found. Please build the environment first.")
        else:
            await self._shutdown_devbox()
            self._devbox = None

        if self._client:
            try:
                await self._client.aclose()
            except Exception as e:
                self.logger.warning(
                    "Failed to close Runloop client cleanly: %s", e, exc_info=True
                )
        self._client = 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._devbox or not self._client:
            raise RuntimeError("Devbox not found. Please build the environment first.")

        source_path = Path(source_path)

        # Use the OO file interface
        await self._devbox.file.upload(
            path=target_path,
            file=source_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._devbox or not self._client:
            raise RuntimeError("Devbox not found. Please build the environment first.")

        source_dir = Path(source_dir)

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

                # Create parent directory if needed
                parent_dir = str(Path(destination_path).parent)
                if parent_dir != ".":
                    await self.exec(f"mkdir -p {parent_dir}")

                await self.upload_file(file_path, destination_path)

    @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._devbox or not self._client:
            raise RuntimeError("Devbox not found. Please build the environment first.")

        target_path = Path(target_path)
        target_path.parent.mkdir(parents=True, exist_ok=True)

        # Use the OO file interface to download bytes directly
        data = await self._devbox.file.download(path=source_path)
        target_path.write_bytes(data)

    @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._devbox or not self._client:
            raise RuntimeError("Devbox not found. Please build the environment first.")

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

        # List all files in the source directory
        list_result = await self.exec(f"find {source_dir} -type f")

        if list_result.stdout:
            files = list_result.stdout.strip().split("\n")

            for file_path in files:
                file_path = file_path.strip()
                if file_path:
                    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)
                    await self.download_file(file_path, local_file_path)

    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._devbox or not self._client:
            raise RuntimeError("Devbox not found. Please build the environment first.")

        # Wrap command with bash -ic
        # full_command = f"bash -ic {shlex.quote(command)}"
        full_command = command

        # Add environment variables
        if env:
            for key, value in env.items():
                full_command = f"{key}={shlex.quote(value)} {full_command}"

        # Add working directory
        if cwd:
            full_command = f"cd {cwd} && {full_command}"
        elif self._workdir:
            full_command = f"cd {self._workdir} && {full_command}"

        try:
            # Execute the command and await completion
            result = await self._devbox.cmd.exec(
                command=full_command,
                shell_name=self._shell_name,
                timeout=timeout_sec,
            )

            stdout_text = await result.stdout()
            stderr_text = await result.stderr()

            exit_code = result.exit_code
            if exit_code is None:
                exit_code = 1  # non-zero to indicate failure

            return ExecResult(
                stdout=stdout_text,
                stderr=stderr_text,
                return_code=exit_code,
            )
        except Exception as e:
            # Ensure we always return an ExecResult, even on errors such as timeouts.
            self.logger.warning(
                "Runloop exec failed for command %r: %s", command, e, exc_info=True
            )
            return ExecResult(
                stdout="",
                stderr=str(e),
                return_code=1,
            )
