from pathlib import Path

import modal
from modal import App, Image, Sandbox, Secret

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

modal.enable_output()


class ModalEnvironment(BaseEnvironment):
    environment_dir: Path
    environment_name: str
    session_id: str
    trial_paths: TrialPaths
    config: EnvironmentConfig

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

    @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."
            )

    def __init__(
        self,
        environment_dir: Path,
        environment_name: str,
        session_id: str,
        trial_paths: TrialPaths,
        task_env_config: EnvironmentConfig,
        add_python_version: str | None = None,
        *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._image: Image | None = None
        self._app: App | None = None
        self._sandbox: Sandbox | None = None
        self._add_python_version = add_python_version

    async def start(self, force_build: bool) -> None:
        """Starts the environment and optionally forces a build."""

        if not self._add_python_version:
            self.logger.warning("Python must be installed in the environment image.")

        self._image = Image.from_dockerfile(
            self._environment_definition_path, add_python=self._add_python_version
        )

        self._app = await App.lookup.aio(self.environment_name, create_if_missing=True)

        self._sandbox = await Sandbox.create.aio(
            app=self._app,
            image=self._image,
            timeout=60 * 60 * 24,
            name=self.session_id,
            cpu=self.task_env_config.cpus,
            memory=self.task_env_config.memory_mb,
        )

        await self._sandbox.mkdir.aio(str(EnvironmentPaths.agent_dir), parents=True)
        await self._sandbox.mkdir.aio(str(EnvironmentPaths.verifier_dir), parents=True)

    async def stop(self, delete: bool):
        """Stops the environment and optionally deletes it."""
        if not self._sandbox:
            return

        await self._sandbox.terminate.aio()
        await self._sandbox.wait.aio(raise_on_termination=False)

    async def upload_file(self, source_path: Path | str, target_path: str):
        """
        Adds a local file to the environment.

        Args:
            source_path: The path to the source local file.
            target_path: The path to which to copy the file.
        """
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please start the environment first.")

        async with await self._sandbox.open.aio(target_path, "wb") as file_handle:
            with open(source_path, "rb") as local_file:
                while True:
                    chunk = local_file.read(8192)
                    if not chunk:
                        break
                    file_handle.write(chunk)

    async def upload_dir(self, source_dir: Path | str, target_dir: str):
        """
        Adds a local directory to the environment.

        Args:
            source_dir: The path to the source local directory.
            target_dir: The path to which to copy the directory.
        """
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please start the environment first.")

        source_path = Path(source_dir)
        if not source_path.exists():
            raise FileNotFoundError(f"Source directory {source_dir} does not exist")

        await self.exec(f"mkdir -p {target_dir}")

        for file_path in source_path.rglob("*"):
            if file_path.is_file():
                relative_path = file_path.relative_to(source_path)
                target_file_path = str(Path(target_dir) / relative_path)
                await self.upload_file(file_path, target_file_path)

    async def download_file(self, source_path: str, target_path: Path | str):
        """
        Downloads a file from the environment to the local machine.

        Args:
            source_path: The path to the source file in the environment.
            target_path: The local path to which to copy the file.
        """
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please start the environment first.")

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

        async with await self._sandbox.open.aio(source_path, "rb") as file_handle:
            with open(target_path, "wb") as local_file:
                while True:
                    chunk = file_handle.read(8192)
                    if not chunk:
                        break
                    local_file.write(chunk)

    async def download_dir(self, source_dir: str, target_dir: Path | str):
        """
        Downloads a directory from the environment to the local machine. This overwrites
        existing files in the target directory.

        Args:
            source_dir: The path to the source directory in the environment.
            target_dir: The local path to which to copy the directory.
        """
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please start the environment first.")

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

        children = await self._sandbox.ls.aio(source_dir)

        for child_name in children:
            # Modal's ls returns just filenames, not full paths
            # Construct the full path by joining source_dir with the child name
            child_path = str(Path(source_dir) / child_name)
            relative_path = Path(child_name)
            local_path = target_dir / relative_path

            try:
                await self._sandbox.ls.aio(child_path)
                await self.download_dir(child_path, local_path)
            except NotADirectoryError:
                local_path.parent.mkdir(parents=True, exist_ok=True)
                await self.download_file(child_path, local_path)

    async def exec(
        self,
        command: str,
        cwd: str | None = None,
        env: dict[str, str] | None = None,
        timeout_sec: int | None = None,
    ) -> ExecResult:
        """
        Executes a command in the environment.

        Args:
            command: The command to execute.
            cwd: The working directory in which to execute the command.
            env: The environment  variables to set.
            timeout_sec: The timeout in seconds.
        """
        if not self._sandbox:
            raise RuntimeError("Sandbox not found. Please start the environment first.")

        process = await self._sandbox.exec.aio(
            "bash",
            "-ic",
            command,
            workdir=cwd,
            secrets=[Secret.from_dict(env)] if env else [],  # type: ignore
            timeout=timeout_sec,
        )

        return_code = await process.wait.aio()

        stdout = await process.stdout.read.aio()
        stderr = await process.stderr.read.aio()

        return ExecResult(
            stdout=stdout,
            stderr=stderr,
            return_code=return_code,
        )
