import asyncio
import json
from pathlib import Path

from dirhash import dirhash
from e2b import AsyncSandbox, AsyncTemplate, FileType, Template
from e2b.sandbox.filesystem.filesystem import WriteEntry
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 E2BEnvironment(BaseEnvironment):
    _UPLOAD_BATCH_SIZE = 20

    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) find a better way to do this

        self._sandbox: AsyncSandbox | None = None
        self._template_name = f"{environment_name}__{dirhash(self.environment_dir, 'sha256')[:8]}".replace(
            ".", "-"
        )

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

    @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_template(self):
        if self.task_env_config.docker_image:
            template = Template().from_image(
                image=self.task_env_config.docker_image,
            )
        else:
            template = Template().from_dockerfile(
                dockerfile_content_or_path=str(self._environment_definition_path),
            )

        await AsyncTemplate.build(
            template=template,
            alias=self._template_name,
            cpu_count=self.task_env_config.cpus,
            memory_mb=self.task_env_config.memory_mb,
        )

    @retry(
        stop=stop_after_attempt(2),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def _create_sandbox(self):
        metadata = {
            "environment_name": self.environment_name,
            "session_id": self.session_id,
        }

        self._sandbox = await AsyncSandbox.create(
            template=self._template_name,
            metadata=metadata,
            timeout=86_400,
        )

    async def _does_template_exist(self) -> bool:
        process = await asyncio.create_subprocess_exec(
            "e2b",
            "template",
            "list",
            "--format",
            "json",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await process.communicate()

        templates = json.loads(stdout.decode("utf-8", errors="ignore").strip())

        return self._template_name in {
            template["aliases"][0]
            for template in templates
            if template["buildStatus"] != "failed" and len(template["aliases"]) > 0
        }

    async def start(self, force_build: bool):
        if force_build or not await self._does_template_exist():
            self.logger.debug(f"Creating template {self._template_name}")

            await self._create_template()

        await self._create_sandbox()

        if not self._sandbox:
            raise RuntimeError(
                "Sandbox not found but was just created. This should never happen."
            )

        await self._sandbox.files.make_dir(str(EnvironmentPaths.agent_dir))
        await self._sandbox.files.make_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.kill()

    async def stop(self, delete: bool):
        """Stops the environment and optionally deletes it."""
        if not delete:
            self.logger.info(
                "E2B harbor are ephemeral and will be deleted after use, "
                "regardless of delete=False."
            )

        if self._sandbox:
            try:
                await self._stop_sandbox()
            except Exception as e:
                self.logger.error(f"Error stopping sandbox: {e}")
            finally:
                self._sandbox = None
        else:
            self.logger.info("Sandbox has already been removed.")

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

        await self._sandbox.files.write(target_path, Path(source_path).read_bytes())

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

        files: list[WriteEntry] = []
        for file_path in Path(source_dir).rglob("*"):
            if file_path.is_file():
                files.append(
                    WriteEntry(
                        path=str(
                            Path(target_dir) / file_path.relative_to(Path(source_dir))
                        ),
                        data=file_path.read_bytes(),
                    )
                )

        if files:
            for i in range(0, len(files), self._UPLOAD_BATCH_SIZE):
                batch = files[i : i + self._UPLOAD_BATCH_SIZE]
                await self._sandbox.files.write_files(batch)

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

        Path(target_path).write_bytes(
            await self._sandbox.files.read(source_path, format="bytes"),
        )

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

        results = await self._sandbox.files.list(source_dir)

        for result in results:
            if result.type == FileType.DIR:
                sub_target_dir = Path(target_dir) / Path(result.path).relative_to(
                    Path(source_dir)
                )
                sub_target_dir.mkdir(parents=True, exist_ok=True)

                await self.download_dir(
                    source_dir=result.path,
                    target_dir=sub_target_dir,
                )

            if result.type == FileType.FILE:
                target_path = Path(target_dir) / Path(result.path).relative_to(
                    Path(source_dir)
                )

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

                await self.download_file(
                    source_path=result.path,
                    target_path=str(target_path),
                )

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    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.")

        handle = await self._sandbox.commands.run(
            cmd=command,
            background=True,
            cwd=cwd or self._workdir,
            envs=env,
            timeout=timeout_sec or 0,
            user="root",
        )

        result = await handle.wait()

        return ExecResult(
            stdout=result.stdout,
            stderr=result.stderr,
            return_code=result.exit_code,
        )
