# (c) 2015-2023 Acellera Ltd http://www.acellera.com
# All Rights Reserved
# Distributed under HTMD Software License Agreement
# No redistribution in whole or part
#
"""Docker/GCloud manifest backend."""

import base64
import gzip
import json
import logging
import os
import subprocess
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Optional, Tuple
from playmolecule._config import _get_config

logger = logging.getLogger("playmolecule.backends.docker")

# Cached credentials and session
_cached_credentials = None
_cached_session = None

_CONFIG = _get_config()


def _image_tag_to_name(image_tag: str) -> str:
    """Convert an image tag to a name."""
    image_name = image_tag.split("/")[-2:]
    image_name = "_".join(image_name)
    image_name = image_name.replace(":", "_")
    return image_name


def _get_or_create_session():
    """Get or create a cached requests session with authentication."""
    import google.auth
    from google.auth.transport.requests import Request
    import requests

    global _cached_credentials, _cached_session

    # Check if we need to refresh credentials
    need_refresh = True
    if _cached_credentials is not None:
        expiry = getattr(_cached_credentials, "expiry", None)
        if expiry is not None:
            if expiry.tzinfo is None:
                expiry = expiry.replace(tzinfo=timezone.utc)
            if expiry > datetime.now(timezone.utc):
                need_refresh = False

    if need_refresh:
        credentials, _ = google.auth.default(
            scopes=["https://www.googleapis.com/auth/cloud-platform"]
        )
        credentials.refresh(Request())
        _cached_credentials = credentials

    if _cached_session is None:
        _cached_session = requests.Session()

    return _cached_session, _cached_credentials.token


def _get_location_and_project_id(root: str) -> Tuple[str, str]:
    # Given a URL like this docker://europe-southwest1-docker.pkg.dev/repositories-368911
    # parse it to get the project_id and location. It could also end with a trailing slash.
    import re

    match = re.match(r"docker://(.*)-docker\.pkg\.dev/(.*)/?", root)
    if match:
        return match.group(1), match.group(2)
    else:
        raise ValueError(f"Invalid Docker URL: {root}")


class _DockerManifestBackend:
    """Backend for loading app manifests from Docker/GCloud registry."""

    def __init__(self, root: str):
        """Initialize the Docker manifest backend.

        Parameters
        ----------
        root : str
            Root path (should start with "docker")
        """
        self.root = root
        self.location, self.project_id = _get_location_and_project_id(root)

    def get_apps(self) -> Dict[str, Dict[str, dict]]:
        """Get all available apps from Docker registry.

        Returns
        -------
        dict
            Nested dict: {appname: {version: {"manifest": ..., "appdir": None, "run.sh": ...}}}
        """
        prefix = f"{self.location}-docker.pkg.dev/{self.project_id}"
        manifests = self._get_app_manifests(prefix)

        apps = {}
        for image_tag, manifest in manifests.items():
            container_config = manifest["container_config"]
            app_name = container_config["name"].lower()
            version = container_config["version"]

            # Generate run script
            run_sh = self._generate_run_script(image_tag)

            if app_name not in apps:
                apps[app_name] = {}
            apps[app_name][f"v{version}"] = {
                "manifest": manifest,
                "appdir": None,
                "run.sh": run_sh,
                "container_image": image_tag,
                "files": self._get_app_files(manifest, image_tag),
            }

        return apps

    def _get_app_files(self, manifest: Optional[dict], container_image: str) -> dict:
        """Get files associated with an app.

        Parameters
        ----------
        manifest : dict or None
            The app manifest
        container_image : str
            The container image tag

        Returns
        -------
        dict
            Dictionary of filename -> _DockerFile objects
        """
        from playmolecule._appfiles import _DockerFile

        if manifest is None:
            return {}

        files = {}
        for name, fullpath in manifest.get("files", {}).items():
            # Add parent directories
            for dir_name, dir_path in self._get_parent_dirs(name, fullpath):
                if dir_name not in files:
                    files[dir_name] = _DockerFile(dir_name, dir_path, container_image)
            files[name] = _DockerFile(name, fullpath, container_image)

        return files

    def _get_parent_dirs(self, name: str, fullpath: str):
        """Yield parent directories for a file path."""
        fullpath_p = Path(fullpath)
        name_p = Path(name)

        for i in range(1, len(name_p.parts)):
            yield str(Path(*name_p.parts[:i])), str(
                Path(*fullpath_p.parts[: -len(name_p.parts) + i])
            )

    def _generate_run_script(self, image_tag: str) -> str:
        """Generate the docker run script for an image."""
        curr_dir = os.path.dirname(os.path.abspath(__file__))
        share_dir = os.path.join(os.path.dirname(curr_dir), "share")

        if _CONFIG.runtime == "docker":
            with open(os.path.join(share_dir, "docker_run.sh"), "r") as f:
                run_sh = f.read()
        elif _CONFIG.runtime == "apptainer":
            with open(os.path.join(share_dir, "apptainer_run.sh"), "r") as f:
                run_sh = f.read()
            run_sh = run_sh.replace("{license_file_or_server}", '"27000@127.0.0.1"')
            run_sh = run_sh.replace("{sif_cache_dir}", _CONFIG.sif_cache_dir)
            run_sh = run_sh.replace("{image_name}", _image_tag_to_name(image_tag))
        else:
            raise RuntimeError(f"Unsupported runtime: {_CONFIG.runtime}")

        return run_sh.replace("{docker_image}", f'"{image_tag}"')

    def _get_app_list(self) -> list:
        """Get list of all Docker images from GCloud Artifact Registry."""
        from google.cloud import artifactregistry_v1

        ar_client = artifactregistry_v1.ArtifactRegistryClient()
        parent = f"projects/{self.project_id.strip()}/locations/{self.location}"
        registry_host = f"{self.location}-docker.pkg.dev"

        repos_request = artifactregistry_v1.ListRepositoriesRequest(parent=parent)
        repos = list(ar_client.list_repositories(request=repos_request))

        def process_repo(repo):
            if repo.format_ != artifactregistry_v1.Repository.Format.DOCKER:
                return []

            repo_id = repo.name.split("/")[-1]
            if repo_id == "acellera-docker-apps":
                return []

            try:
                pkg_request = artifactregistry_v1.ListPackagesRequest(parent=repo.name)
                packages = ar_client.list_packages(request=pkg_request)
            except Exception:
                # If we don't have permissions to list packages in that repo, return empty list
                return []

            repo_images = []
            for pkg in packages:
                image_name = pkg.name.split("/")[-1]
                full_tag = (
                    f"{registry_host}/{self.project_id}/{repo_id}/{image_name}:latest"
                )

                repo_images.append(
                    {
                        "project_id": self.project_id,
                        "location": self.location,
                        "repo": repo_id,
                        "image": image_name,
                        "tag": "latest",
                        "full_tag": full_tag,
                    }
                )
            return repo_images

        all_images = []
        with ThreadPoolExecutor(max_workers=20) as executor:
            results = executor.map(process_repo, repos)
            for result in results:
                all_images.extend(result)

        return all_images

    def _get_local_labels(self, image_tag: str) -> Optional[dict]:
        """Get all labels from the local image."""

        if _CONFIG.runtime == "docker":
            import docker

            client = docker.from_env()
            image = client.images.get(image_tag)
            return image.labels
        elif _CONFIG.runtime == "apptainer":
            # For apptainer, we need to check if we have a cached SIF file
            # and inspect it for labels
            sif_path = os.path.join(
                _CONFIG.sif_cache_dir, _image_tag_to_name(image_tag) + ".sif"
            )

            if not os.path.exists(sif_path):
                # Image not cached locally
                raise ValueError(f"Image not found locally: {image_tag}")

            # Use apptainer inspect to get labels
            try:
                result = subprocess.run(
                    ["apptainer", "inspect", "--labels", "--json", sif_path],
                    capture_output=True,
                    text=True,
                    timeout=30,
                )
                if result.returncode == 0:
                    labels_data = json.loads(result.stdout)
                    return (
                        labels_data.get("data", {})
                        .get("attributes", {})
                        .get("labels", {})
                    )
                else:
                    logger.error(f"Failed to inspect image {sif_path}: {result.stderr}")
                    return None
            except Exception as e:
                logger.error(f"Error inspecting apptainer image: {e}")
                return None
        else:
            raise RuntimeError(f"Unsupported runtime: {_CONFIG.runtime}")

    def _get_local_digest(self, image_tag: str) -> Optional[str]:
        """Get the digest of a local image."""
        if _CONFIG.runtime == "docker":
            import docker

            try:
                client = docker.from_env()
                image = client.images.get(image_tag)
                # Docker stores the digest in RepoDigests
                repo_digests = image.attrs.get("RepoDigests", [])
                for digest in repo_digests:
                    if image_tag in digest:
                        # Extract the digest part after @
                        return digest.split("@")[-1]
                return None
            except Exception:
                return None

        elif _CONFIG.runtime == "apptainer":
            # For apptainer, we need to store the digest when we pull the image
            # and retrieve it from a metadata file
            sif_path = os.path.join(
                _CONFIG.sif_cache_dir, _image_tag_to_name(image_tag) + ".sif"
            )
            digest_file = sif_path + ".digest"

            if os.path.exists(digest_file):
                try:
                    with open(digest_file, "r") as f:
                        return f.read().strip()
                except Exception:
                    return None
            return None
        else:
            raise RuntimeError(f"Unsupported runtime: {_CONFIG.runtime}")

    def _get_remote_labels(
        self,
        project_id: str,
        location: str,
        repo: str,
        image: str,
        tag: str,
    ) -> Optional[dict]:
        """Fetch all labels from the remote image."""
        t_start = time.time()
        session, token = _get_or_create_session()
        logger.debug(f"  Time to get session/token: {time.time() - t_start:.3f}s")

        base_url = f"https://{location}-docker.pkg.dev"
        image_path = f"v2/{project_id}/{repo}/{image}"

        headers = {
            "Authorization": f"Bearer {token}",
            "Accept": (
                "application/vnd.docker.distribution.manifest.v2+json, "
                "application/vnd.oci.image.index.v1+json, "
                "application/vnd.oci.image.manifest.v1+json"
            ),
        }

        # Get manifest
        manifest_url = f"{base_url}/{image_path}/manifests/{tag}"
        t = time.time()
        r = session.get(manifest_url, headers=headers)
        if r.status_code != 200:
            raise ValueError(f"Failed to get manifest: {r.status_code} {r.text}")
        logger.debug(f"  Time to get manifest: {time.time() - t:.3f}s")

        manifest = r.json()
        media_type = r.headers.get("Content-Type", "")

        # If it's an Index, pick the first manifest in the list (usually the default)
        if "index" in media_type or "manifests" in manifest:
            # Dig deeper to get the specific manifest for the first architecture
            first_manifest_digest = manifest["manifests"][0]["digest"]

            # Fetch the actual manifest for that specific digest
            manifest_url = f"{base_url}/{image_path}/manifests/{first_manifest_digest}"
            r = session.get(manifest_url, headers=headers)
            manifest = r.json()

        # Extract config digest
        config_digest = manifest.get("config", {}).get("digest")
        if not config_digest:
            config_digest = manifest.get("manifests", [{}])[0].get("digest")
        if not config_digest:
            raise ValueError("Could not find config digest in manifest")

        # Get config blob
        blob_url = f"{base_url}/{image_path}/blobs/{config_digest}"
        t = time.time()
        r = session.get(blob_url, headers=headers)
        if r.status_code != 200:
            raise ValueError(f"Failed to get config blob: {r.status_code}")
        logger.debug(f"  Time to get config blob: {time.time() - t:.3f}s")

        config_data = r.json()

        # Extract and decode label
        labels = config_data.get("config", {}).get("Labels", {})
        return labels

    def _get_remote_digest(
        self,
        project_id: str,
        location: str,
        repo: str,
        image: str,
        tag: str,
    ) -> Optional[str]:
        """Get the digest of a remote image without downloading it."""
        session, token = _get_or_create_session()

        base_url = f"https://{location}-docker.pkg.dev"
        image_path = f"v2/{project_id}/{repo}/{image}"

        headers = {
            "Authorization": f"Bearer {token}",
            "Accept": (
                "application/vnd.docker.distribution.manifest.v2+json, "
                "application/vnd.oci.image.index.v1+json, "
                "application/vnd.oci.image.manifest.v1+json"
            ),
        }

        # Use HEAD request to get digest without downloading manifest body
        manifest_url = f"{base_url}/{image_path}/manifests/{tag}"
        r = session.head(manifest_url, headers=headers)

        if r.status_code != 200:
            # Fallback to GET if HEAD is not supported
            r = session.get(manifest_url, headers=headers)
            if r.status_code != 200:
                raise ValueError(f"Failed to get manifest digest: {r.status_code}")

        # The digest is in the Docker-Content-Digest header
        digest = r.headers.get("Docker-Content-Digest")
        if not digest:
            # For OCI manifests, it might be in a different header
            digest = r.headers.get("Digest")

        return digest

    def _get_app_manifests(self, prefix: str) -> Dict[str, dict]:
        """Get manifests for all images matching prefix."""

        manifest_label = "function.manifest.compressed"
        docs_label = "function.docs.compressed"

        images = self._get_app_list()

        # Filter by prefix
        target_images = [img for img in images if img["full_tag"].startswith(prefix)]

        labels = {}

        # Separate local and remote images
        local_images = []
        remote_images = []

        if _CONFIG.runtime == "docker":
            import docker

            client = docker.from_env()

            for image_info in target_images:
                image_tag = image_info["full_tag"]
                try:
                    client.images.get(image_tag)
                    local_images.append(image_info)
                except docker.errors.ImageNotFound:
                    remote_images.append(image_info)
        elif _CONFIG.runtime == "apptainer":
            # Check for cached SIF files
            for image_info in target_images:
                image_tag = image_info["full_tag"]
                sif_path = os.path.join(
                    _CONFIG.sif_cache_dir, _image_tag_to_name(image_tag) + ".sif"
                )

                if os.path.exists(sif_path):
                    local_images.append(image_info)
                else:
                    remote_images.append(image_info)
        else:
            # No runtime available, all images are remote
            remote_images = target_images

        # Process local images
        for image_info in local_images:
            image_tag = image_info["full_tag"]

            try:
                labels[image_tag] = self._get_local_labels(image_tag)
                logger.info(f"  {image_tag} (local)")
            except Exception as e:
                logger.error(f"  Failed to get labels for {image_tag}: {e}")
                continue

        # Process remote images in parallel
        if remote_images:
            t_start = time.time()

            def fetch_remote_labels(image_info):
                try:
                    labels = self._get_remote_labels(
                        image_info["project_id"],
                        image_info["location"],
                        image_info["repo"],
                        image_info["image"],
                        image_info["tag"],
                    )
                    return image_info["full_tag"], labels, None
                except Exception as e:
                    return image_info["full_tag"], None, str(e)

            with ThreadPoolExecutor(max_workers=10) as executor:
                future_to_image = {
                    executor.submit(fetch_remote_labels, img): img
                    for img in remote_images
                }

                for future in as_completed(future_to_image):
                    image_tag, image_labels, error = future.result()

                    if error:
                        logger.error(f"  Error fetching {image_tag}: {error}")
                    elif image_labels:
                        labels[image_tag] = image_labels
                        logger.info(f"  {image_tag}")

            elapsed = time.time() - t_start
            logger.info(
                f"Loaded {len(remote_images)} app manifests in {elapsed:.2f}s "
                f"({elapsed/len(remote_images):.2f}s per app)"
            )

        manifests = {}
        for image_tag, labels in labels.items():
            try:
                manifest = labels.get(manifest_label)
                if manifest:
                    manifests[image_tag] = json.loads(
                        gzip.decompress(base64.b64decode(manifest))
                    )
            except Exception as e:
                logger.error(f"Failed to decompress manifest for {image_tag}: {e}")
                continue
            try:
                docs = labels.get(docs_label)
                if docs:
                    manifests[image_tag]["docs"] = json.loads(
                        gzip.decompress(base64.b64decode(docs))
                    )
            except Exception as e:
                logger.error(f"Failed to decompress docs for {image_tag}: {e}")
                continue

        return manifests

    def _save_digest_for_apptainer(self, image_tag: str, digest: str) -> None:
        """Save the digest for an Apptainer image."""
        if _CONFIG.runtime == "apptainer":
            sif_path = os.path.join(
                _CONFIG.sif_cache_dir, _image_tag_to_name(image_tag) + ".sif"
            )
            digest_file = sif_path + ".digest"
            try:
                with open(digest_file, "w") as f:
                    f.write(digest)
            except Exception as e:
                logger.warning(f"Failed to save digest for {image_tag}: {e}")


def _update_docker_apps_from_gcloud(
    pull_new: bool = False, interactive: bool = False
) -> None:
    """Pull updates for container images from GCloud Artifact Registry.

    Parameters
    ----------
    pull_new : bool, optional
        If True, also pull images that don't exist locally. Default is False (only update existing).
    interactive : bool, optional
        If True, show a list of all apps and allow user to select which ones to update.
        Default is False (update all existing or all if pull_new is True).
    """
    # Find the docker registry in the registries list
    docker_registry = [
        registry for registry in _CONFIG.registries if registry.startswith("docker")
    ]
    if len(docker_registry) == 0:
        raise ValueError("No docker registry found in registries list")
    docker_registry = docker_registry[0]

    backend = _DockerManifestBackend(docker_registry)
    images = backend._get_app_list()

    # Interactive selection
    if interactive:
        # In interactive mode, always enable pull_new since the user is explicitly choosing apps
        pull_new = True
        # Collect information about which images exist locally
        local_images = set()
        if _CONFIG.runtime == "docker":
            import docker

            docker_client = docker.from_env()
            for image_info in images:
                tag = image_info["full_tag"]
                try:
                    docker_client.images.get(tag)
                    local_images.add(tag)
                except docker.errors.ImageNotFound:
                    pass
        elif _CONFIG.runtime == "apptainer":
            os.makedirs(_CONFIG.sif_cache_dir, exist_ok=True)
            for image_info in images:
                tag = image_info["full_tag"]
                image_name = _image_tag_to_name(tag)
                sif_path = os.path.join(_CONFIG.sif_cache_dir, image_name + ".sif")
                if os.path.exists(sif_path):
                    local_images.add(tag)

        # Display available images and their status
        print("\nAvailable PlayMolecule Apps:")
        print("-" * 80)
        print(
            f"{'#':<4} {'App Name':<25} {'Container Ver':<15} {'Tag':<10} {'Status':<15}"
        )
        print("-" * 80)
        for i, image_info in enumerate(images):
            tag = image_info["full_tag"]
            repo = image_info.get("repo", "")
            app_name = tag.split("/")[-1].split(":")[0]

            # Extract container version from repo (e.g., "parameterize-v1" -> "v1")
            container_version = ""
            if "-v" in repo:
                parts = repo.split("-v")
                if len(parts) >= 2 and parts[-1].replace(".", "").isdigit():
                    container_version = f"v{parts[-1]}"

            tag_version = tag.split(":")[-1] if ":" in tag else "latest"
            status = "✓ Installed" if tag in local_images else "✗ Not installed"
            print(
                f"{i+1:<4} {app_name:<25} {container_version:<15} {tag_version:<10} {status:<15}"
            )

        print("\nOptions:")
        print(
            "  - Enter numbers separated by spaces to select specific apps (e.g., '1 3 5')"
        )
        print("  - Enter 'all' to select all apps")
        print("  - Enter 'installed' to select only installed apps")
        print("  - Enter 'new' to select only non-installed apps")
        print("  - Press Enter to cancel")

        selection = input("\nYour selection: ").strip().lower()

        if not selection:
            logger.info("Update cancelled by user")
            return

        selected_indices = []
        if selection == "all":
            selected_indices = list(range(len(images)))
        elif selection == "installed":
            selected_indices = [
                i for i, img in enumerate(images) if img["full_tag"] in local_images
            ]
        elif selection == "new":
            selected_indices = [
                i for i, img in enumerate(images) if img["full_tag"] not in local_images
            ]
        else:
            try:
                selected_indices = [int(x) - 1 for x in selection.split()]
                # Validate indices
                for idx in selected_indices:
                    if idx < 0 or idx >= len(images):
                        raise ValueError(f"Invalid selection: {idx + 1}")
            except ValueError as e:
                logger.error(f"Invalid selection: {e}")
                return

        # Filter images to only selected ones
        images = [images[i] for i in selected_indices]

        if not images:
            logger.info("No apps selected for update")
            return

        print(f"\nSelected {len(images)} app(s) for update")

    if _CONFIG.runtime == "docker":
        import docker

        docker_client = docker.from_env()

        # Process images based on pull_new parameter
        for image_info in images:
            tag = image_info["full_tag"]
            try:
                # Check if image exists locally
                docker_client.images.get(tag)
                # If we get here, the image exists locally, so update it
                logger.info(f"Updating existing image: {tag} ...")
                docker_client.images.pull(tag)
            except docker.errors.ImageNotFound:
                # Image doesn't exist locally
                if pull_new:
                    logger.info(f"Pulling new image: {tag} ...")
                    try:
                        docker_client.images.pull(tag)
                    except Exception as e:
                        logger.error(f"Failed to pull {tag}: {e}")
                else:
                    logger.debug(f"Skipping {tag} - not found locally")
            except Exception as e:
                logger.error(f"Failed to pull {tag}: {e}")
    elif _CONFIG.runtime == "apptainer":
        os.makedirs(_CONFIG.sif_cache_dir, exist_ok=True)

        for image_info in images:
            tag = image_info["full_tag"]
            image_name = _image_tag_to_name(tag)
            sif_path = os.path.join(_CONFIG.sif_cache_dir, image_name + ".sif")

            if os.path.exists(sif_path):
                # Image exists locally, check if update is needed
                try:
                    # Get local and remote digests
                    local_digest = backend._get_local_digest(tag)
                    remote_digest = backend._get_remote_digest(
                        image_info["project_id"],
                        image_info["location"],
                        image_info["repo"],
                        image_info["image"],
                        image_info["tag"],
                    )

                    if local_digest and remote_digest and local_digest == remote_digest:
                        logger.info(f"Skipping {tag} - already up to date")
                        continue

                    logger.info(f"Updating existing image: {tag} ...")

                    # Pull to a temporary file first
                    temp_sif_path = sif_path + ".tmp"

                    result = subprocess.run(
                        [
                            "apptainer",
                            "pull",
                            "--force",
                            temp_sif_path,
                            f"docker://{tag}",
                        ],
                        capture_output=True,
                        text=True,
                        timeout=300,  # 5 minute timeout
                    )
                    if result.returncode != 0:
                        logger.error(f"Failed to pull {tag}: {result.stderr}")
                        # Clean up temp file if it exists
                        if os.path.exists(temp_sif_path):
                            os.remove(temp_sif_path)
                    else:
                        # Successfully downloaded, now replace the old file
                        try:
                            # Move temp file to final location (atomic on same filesystem)
                            os.replace(temp_sif_path, sif_path)

                            # Update the digest file
                            if remote_digest:
                                backend._save_digest_for_apptainer(tag, remote_digest)

                            logger.info(f"Successfully updated {tag}")
                        except Exception as e:
                            logger.error(f"Failed to replace old image: {e}")
                            # Clean up temp file
                            if os.path.exists(temp_sif_path):
                                os.remove(temp_sif_path)

                except Exception as e:
                    logger.error(f"Failed to check/update {tag}: {e}")
            else:
                # Image doesn't exist locally
                if pull_new:
                    logger.info(f"Pulling new image: {tag} ...")
                    try:
                        # Get remote digest first
                        remote_digest = backend._get_remote_digest(
                            image_info["project_id"],
                            image_info["location"],
                            image_info["repo"],
                            image_info["image"],
                            image_info["tag"],
                        )

                        result = subprocess.run(
                            [
                                "apptainer",
                                "pull",
                                "--force",
                                sif_path,
                                f"docker://{tag}",
                            ],
                            capture_output=True,
                            text=True,
                            timeout=300,  # 5 minute timeout
                        )
                        if result.returncode != 0:
                            logger.error(f"Failed to pull {tag}: {result.stderr}")
                        else:
                            # Save the digest for future comparisons
                            if remote_digest:
                                backend._save_digest_for_apptainer(tag, remote_digest)
                            logger.info(f"Successfully pulled {tag}")
                    except Exception as e:
                        logger.error(f"Failed to pull {tag}: {e}")
                else:
                    logger.debug(f"Skipping {tag} - not found locally")
    else:
        raise RuntimeError(f"Unsupported runtime: {_CONFIG.runtime}")
