# (c) 2015-2023 Acellera Ltd http://www.acellera.com
# All Rights Reserved
# Distributed under HTMD Software License Agreement
# No redistribution in whole or part
#
from typing import Dict, List, Optional, TYPE_CHECKING

from playmolecule._config import _get_config

if TYPE_CHECKING:
    import requests


_CONFIG = _get_config()


class _Artifacts:
    """Container for app artifacts/datasets."""

    def __init__(self, artifacts: List[dict], files_dict: Dict[str, "_File"]) -> None:
        import os

        searchpaths = [
            ["datasets"],
            ["artifacts"],
            ["files", "datasets"],
            ["files", "artifacts"],
            [],
        ]

        for ds in artifacts:
            path = None
            for sp in searchpaths:
                path = os.path.join(*sp, ds["path"])
                if path in files_dict:
                    break

            if path is None:
                raise RuntimeError(
                    f"Could not find dataset {ds['name']} at path {path}"
                )

            try:
                if "." in ds["name"]:
                    raise RuntimeError(
                        f"Dataset/artifact names cannot include dots in the name. {ds['name']} contains a dot."
                    )
                if not ds["name"][0].isalpha():
                    raise RuntimeError(
                        f"Dataset/artifact names must start with a letter. {ds['name']} does not."
                    )

                # Copy the original file object
                import copy

                new_file = copy.copy(files_dict[path])

                # Update name and description
                new_file.name = ds["name"]
                new_file.description = ds.get("description", new_file.description)

                setattr(self, ds["name"], new_file)
            except Exception:
                pass

    def __str__(self) -> str:
        descr = ""
        for key in self.__dict__:
            descr += f"{self.__dict__[key]}\n"
        return descr

    def __repr__(self) -> str:
        return self.__str__()


class _File:
    """Represents a file in an app's files directory."""

    def __init__(self, name: str, path: str, description: Optional[str] = None) -> None:
        self.name = name
        self.path = path
        self.description = description

    def __str__(self) -> str:
        string = f"[{self.name}] {self.path}"
        if self.description is not None:
            string += f" '{self.description}'"
        return string

    def __repr__(self) -> str:
        return self.__str__()

    def download(self, output_path: Optional[str] = None) -> str:
        """Download (copy) the file or directory to the specified path.

        Parameters
        ----------
        output_path : str, optional
            Path where to save the file/directory. If None, saves to current
            directory with the same name.

        Returns
        -------
        str
            Path to the downloaded file/directory
        """
        import os
        import shutil

        if output_path is None:
            output_path = os.path.basename(self.path)

        # Check if source exists
        if not os.path.exists(self.path):
            raise FileNotFoundError(f"Source file/directory not found: {self.path}")

        # Get absolute paths to properly compare
        source_abs = os.path.abspath(self.path)
        output_abs = os.path.abspath(output_path)

        # Check if source and destination are the same
        if source_abs == output_abs:
            # Already at the destination, nothing to do
            return output_path

        # Handle directories
        if os.path.isdir(self.path):
            # Remove existing directory if it exists
            if os.path.exists(output_path):
                if os.path.isdir(output_path):
                    shutil.rmtree(output_path)
                else:
                    os.remove(output_path)
            # Copy directory
            shutil.copytree(self.path, output_path)
        else:
            # Handle files
            # Ensure parent directory exists
            output_dir = os.path.dirname(output_path)
            if output_dir:
                os.makedirs(output_dir, exist_ok=True)
            # Copy file
            shutil.copy2(self.path, output_path)

        return output_path


class _DockerFile(_File):
    """Represents a file in a Docker image."""

    def __init__(
        self,
        name: str,
        path: str,
        container_image: str,
        description: Optional[str] = None,
    ) -> None:
        super().__init__(name, path, description)
        self.container_image = container_image

    def download(self, output_path: Optional[str] = None) -> str:
        """Download the file from the Docker image.

        Parameters
        ----------
        output_path : str, optional
            Path where to save the file. If None, saves to current directory
            with the same filename.

        Returns
        -------
        str
            Path to the downloaded file
        """
        import os
        import subprocess
        import tempfile

        if output_path is None:
            output_path = os.path.basename(self.path)

        # Ensure directory exists
        output_dir = os.path.dirname(output_path)
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)

        if _CONFIG.runtime == "docker":
            # Use subprocess for better reliability
            container_name = f"tmp_extract_{os.getpid()}_{id(self)}"
            try:
                # Create container without starting it
                subprocess.run(
                    [
                        "docker",
                        "create",
                        "--name",
                        container_name,
                        self.container_image,
                    ],
                    check=True,
                    capture_output=True,
                )

                # Copy file/directory from container
                subprocess.run(
                    ["docker", "cp", f"{container_name}:{self.path}", output_path],
                    check=True,
                    capture_output=True,
                )
            finally:
                # Clean up container
                subprocess.run(
                    ["docker", "rm", container_name],
                    capture_output=True,
                )

        elif _CONFIG.runtime == "apptainer":
            # For apptainer, we need to extract from the SIF file
            from playmolecule._backends._docker import _image_tag_to_name

            image_name = _image_tag_to_name(self.container_image)
            sif_path = os.path.join(_CONFIG.sif_cache_dir, image_name + ".sif")

            if not os.path.exists(sif_path):
                raise FileNotFoundError(
                    f"SIF file not found at {sif_path}. "
                    f"Please ensure the image {self.container_image} is pulled."
                )

            # Create a temporary directory for extraction
            with tempfile.TemporaryDirectory() as tmp_dir:
                tmp_path = os.path.join(tmp_dir, "extracted")

                # Execute cp command inside the container to copy to bind mount
                # Use cp -r to handle both files and directories
                subprocess.run(
                    [
                        "apptainer",
                        "exec",
                        "--bind",
                        f"{tmp_dir}:/tmp/extract",
                        sif_path,
                        "cp",
                        "-r",
                        self.path,
                        "/tmp/extract/extracted",
                    ],
                    check=True,
                    capture_output=True,
                )

                # Move to final location
                if os.path.exists(tmp_path):
                    # Use shutil.move to handle both files and directories
                    import shutil

                    shutil.move(tmp_path, output_path)
        else:
            raise RuntimeError(
                f"Unsupported runtime: {_CONFIG.runtime}. "
                "Docker or Apptainer runtime required for _DockerFile."
            )

        return output_path


class _HTTPFile(_File):
    """Represents a file in an HTTP server."""

    def __init__(
        self,
        name: str,
        path: str,
        base_url: str,
        session: Optional["requests.Session"] = None,
        description: Optional[str] = None,
    ) -> None:
        super().__init__(name, path, description)
        self.base_url = base_url.rstrip("/")
        self._session = session

    def download(self, output_path: Optional[str] = None) -> str:
        """Download the file from the HTTP server.

        Parameters
        ----------
        output_path : str, optional
            Path where to save the file. If None, saves to current directory
            with the same filename.

        Returns
        -------
        str
            Path to the downloaded file
        """
        import os
        import requests

        if output_path is None:
            output_path = os.path.basename(self.path)

        # Ensure directory exists
        output_dir = os.path.dirname(output_path)
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)

        # Construct full URL
        url = f"{self.base_url}/{self.path.lstrip('/')}"

        # Use provided session or create a new one
        session = self._session or requests.Session()

        try:
            # Download with streaming to handle large files
            response = session.get(url, stream=True, timeout=(10, 300))
            response.raise_for_status()

            # Write to file in chunks
            chunk_size = 8192
            with open(output_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=chunk_size):
                    if chunk:
                        f.write(chunk)
        except requests.exceptions.RequestException as e:
            raise RuntimeError(f"Failed to download file from {url}: {e}")
        finally:
            if not self._session:
                # Close the session if we created it
                session.close()

        return output_path
