import shlex
from pathlib import Path
from typing import Any

from runloop_api_client import AsyncRunloop
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,
            *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
            ),
            "/workspace",
        )

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

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

    @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_devbox(self):
        if not self._client:
            raise RuntimeError("RunLoop client not found. This should never happen.")

        # Create devbox from Dockerfile
        dockerfile_content = self._environment_definition_path.read_text()

        # Create a blueprint with the Dockerfile and wait for it to build
        blueprint = await self._client.blueprints.create_and_await_build_complete(
            dockerfile=dockerfile_content,
            name=f"{self.environment_name}_blueprint",
            launch_parameters={
                "user_parameters": {
                    "username": "root",
                    "uid": 0,
                }
            },
        )

        # Create devbox from the blueprint
        self._devbox = await self._client.devboxes.create_and_await_running(
            blueprint_id=blueprint.id,
            name=self.environment_name,
            launch_parameters={
                "user_parameters": {
                    "username": "root",
                    "uid": 0,
                }
            },
        )

    async def start(self, force_build: bool):
        self._client = AsyncRunloop()

        await self._create_devbox()

        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._client.devboxes.shutdown(self._devbox.id)

    async def stop(self, delete: bool):
        if not self._devbox:
            self.logger.warning("Devbox not found. Please build the environment first.")
        else:
            try:
                await self._shutdown_devbox()
            except Exception as e:
                self.logger.error(f"Error stopping devbox {self._devbox.id}: {e}")
            finally:
                self._devbox = None

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

        await self._client.devboxes.upload_file(
            id=self._devbox.id,
            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)

        response = await self._client.devboxes.download_file(
            id=self._devbox.id,
            path=source_path,
        )

        content = await response.read()
        target_path.write_bytes(content)

    @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}"

        # Execute the command and await completion
        try:
            result = await self._client.devboxes.execute(
                id=self._devbox.id,
                command=full_command,
                shell_name=self._shell_name,
                timeout=timeout_sec,
            )

            return ExecResult(
                stdout=result.stdout or "",
                stderr=result.stderr or "",
                return_code=result.exit_status or 0,
            )
        except Exception as e:
            # If execution fails, return error result
            return ExecResult(
                stdout="",
                stderr=str(e),
                return_code=1,
            )
