# (C) Copyright 2024 Anemoi contributors.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
#
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.


"""Collect information about the current environment, like:

- The Python version
- The versions of the modules which are currently loaded
- The git information for the modules which are currently loaded from a git repository
- ...
"""

import datetime
import json
import logging
import os
import subprocess
import sys
import sysconfig
from functools import cache
from importlib.metadata import PackageNotFoundError
from importlib.metadata import distribution
from typing import Any

LOG = logging.getLogger(__name__)


def lookup_git_repo(path: str) -> Any | None:
    """Lookup the git repository for a given path.

    Parameters
    ----------
    path : str
        The path to lookup.

    Returns
    -------
    Repo, optional
        The git repository if found, otherwise None.
    """
    try:
        from git import InvalidGitRepositoryError
        from git import Repo
    except ImportError:
        return None

    while path != "/":
        try:
            return Repo(path)
        except InvalidGitRepositoryError:
            path = os.path.dirname(path)

    return None


def _check_for_git(paths: list[tuple[str, str]], full: bool) -> dict[str, Any]:
    """Check for git information for the given paths.

    Parameters
    ----------
    paths : list of tuple
        The list of paths to check.
    full : bool
        Whether to collect full information.

    Returns
    -------
    dict
        The git information for the given paths.
    """
    versions = {}
    for name, path in paths:
        repo = lookup_git_repo(path)
        if repo is None:
            continue

        try:

            if not full:
                versions[name] = dict(
                    git=dict(
                        sha1=repo.head.commit.hexsha,
                        modified_files=len([item.a_path for item in repo.index.diff(None)]),
                        untracked_files=len(repo.untracked_files),
                    ),
                )
                continue

            versions[name] = dict(
                path=path,
                git=dict(
                    sha1=repo.head.commit.hexsha,
                    remotes=[r.url for r in repo.remotes],
                    modified_files=sorted([item.a_path for item in repo.index.diff(None)]),
                    untracked_files=sorted(repo.untracked_files),
                ),
            )

        except ValueError as e:
            LOG.error(f"Error checking git repo {path}: {e}")

    return versions


def _package_version(name: str) -> str | None:
    from importlib.metadata import PackageNotFoundError
    from importlib.metadata import version

    # This the recommended way to get the version of a package
    try:
        return version(name)
    except PackageNotFoundError:
        return None


def _get_package_source_url(package_name: str) -> dict[str, Any] | None:
    """Extract the source URL from package metadata if installed from git or other VCS. This reads PEP 610 direct_url.json files created by pip when installing from git URLs.

    Parameters
    ----------
    package_name : str
        The name of the package to check.

    Returns
    -------
    dict or None
        Dictionary with 'url' and optionally 'vcs_info' (commit hash, branch, etc.) if available.
    """
    try:
        dist = distribution(package_name)
    except PackageNotFoundError as e:
        LOG.debug(f"Could not get source URL for {package_name}: {e}")
        return None

    try:
        direct_url_text = dist.read_text("direct_url.json")
    except FileNotFoundError as e:
        LOG.debug(f"No direct_url.json found for {package_name}: {e}")
        return None

    try:
        direct_url = json.loads(str(direct_url_text))
        result = {"url": direct_url.get("url")}
    except json.JSONDecodeError as e:
        LOG.debug(f"Invalid direct_url.json for {package_name}: {e}")
        return None

    # Add VCS info if available (commit hash, requested revision, etc.)
    if "vcs_info" in direct_url:
        result["vcs_info"] = direct_url["vcs_info"]

    # Add subdirectory info if present (e.g., for monorepos)
    if "subdirectory" in direct_url:
        result["subdirectory"] = direct_url["subdirectory"]

    return result


def _module_versions() -> tuple[dict[str, Any], set]:
    """Collect version information for all loaded modules.
       Include source URL information from PEP 610 direct_url.json files.

    Returns
    -------
    tuple of dict and set
        The version information and the set of paths.
    """

    SKIP = set(sys.stdlib_module_names) | set(sys.builtin_module_names)

    paths = set()

    versions = {}

    for name, module in sorted(sys.modules.items()):
        if name in SKIP:
            continue

        version = _package_version(name)
        if version is None:
            continue

        # Store dict with source info
        source_url = _get_package_source_url(name)
        versions[name] = {"version": version}
        if source_url:  # Package contains source info
            versions[name]["source"] = source_url

        if hasattr(module, "__file__") and module.__file__ is not None:
            paths.add((name, os.path.realpath(module.__file__)))

    return versions, paths


@cache
def package_distributions() -> dict[str, list[str]]:
    """Get the package distributions.

    Returns
    -------
    dict
        The package distributions.
    """
    # Takes a significant amount of time to run
    # so cache the result
    from importlib import metadata

    return metadata.packages_distributions()


def import_name_to_distribution_name(packages: list[str]) -> dict[str, str]:
    """Convert import names to distribution names.

    Parameters
    ----------
    packages : list of str
        The list of import names.

    Returns
    -------
    dict
        The dictionary mapping import names to distribution names.
    """

    distribution_names = {}
    package_distribution_names = package_distributions()

    for package in [p for p in packages if p in package_distribution_names]:
        distr_name = package_distribution_names[package]
        if isinstance(distr_name, list):
            if len(distr_name) > 1:
                # Multiple distributions for the same package, i.e. anemoi-graphs, anemoi-utils, ..., Don't know how to handle this
                continue
            distr_name = distr_name[0]

        if distr_name != package:
            distribution_names[package] = distr_name

    return distribution_names


def module_versions(full: bool) -> tuple[dict[str, Any], dict[str, Any]]:
    """Collect version information for all loaded modules and their git information.

    Parameters
    ----------
    full : bool
        Whether to collect full information.

    Returns
    -------
    tuple of dict and dict
        The version information and the git information.
    """
    versions, paths = _module_versions()
    git_versions = _check_for_git(paths, full)
    return versions, git_versions


def _name(obj: Any) -> str:
    """Get the name of an object.

    Parameters
    ----------
    obj : Any
        The object to get the name of.

    Returns
    -------
    str
        The name of the object.
    """
    if hasattr(obj, "__name__"):
        if hasattr(obj, "__module__"):
            return f"{obj.__module__}.{obj.__name__}"
        return obj.__name__
    if hasattr(obj, "__class__"):
        return _name(obj.__class__)
    return str(obj)


def _paths(path_or_object: None | str | list[str] | tuple[str] | Any) -> list[tuple[str, str]]:
    """Get the paths for a given path or object.

    Parameters
    ----------
    path_or_object : str, list, tuple, or object
        The path or object to get the paths for.

    Returns
    -------
    list of tuple
        The list of paths.
    """
    if path_or_object is None:
        _, paths = _module_versions()
        return paths

    if isinstance(path_or_object, (list, tuple, set)):
        paths = []
        for p in path_or_object:
            paths.extend(_paths(p))
        return paths

    if isinstance(path_or_object, str):
        module = sys.modules.get(path_or_object)
        if module is not None:
            return _paths(module)
        return [(path_or_object, path_or_object)]

    if hasattr(path_or_object, "__module__"):
        module = sys.modules.get(path_or_object.__module__)
        return [(path_or_object.__module__, module.__file__)]

    name = _name(path_or_object)
    paths = []
    if hasattr(path_or_object, "__file__"):
        paths.append((name, path_or_object.__file__))

    if hasattr(path_or_object, "__code__"):
        paths.append((name, path_or_object.__code__.co_filename))

    if hasattr(path_or_object, "__module__"):
        module = sys.modules.get(path_or_object.__module__)
        paths.append((name, module.__file__))

    if not paths:
        raise ValueError(f"Could not find path for {name} {path_or_object} {type(path_or_object)}")

    return paths


def git_check(*args: Any) -> dict[str, Any]:
    """Return the git information for the given arguments.

    Arguments can be:
        - an empty list, in that case all loaded modules are checked
        - a module name
        - a module object
        - an object or a class
        - a path to a directory

    Parameters
    ----------
    args : list
        The list of arguments to check

    Returns
    -------
    dict
        An object with the git information for the given arguments.

        >>> {
                "anemoi.utils": {
                    "sha1": "c999d83ae283bcbb99f68d92c42d24315922129f",
                    "remotes": [
                        "git@github.com:ecmwf/anemoi-utils.git"
                    ],
                    "modified_files": [
                        "anemoi/utils/checkpoints.py"
                    ],
                    "untracked_files": []
                }
            }
    """
    paths = _paths(args if len(args) > 0 else None)

    git = _check_for_git(paths, full=True)
    result = {}
    for k, v in git.items():
        result[k] = v["git"]

    return result


def platform_info() -> dict[str, Any]:
    """Get the platform information.

    Returns
    -------
    dict
        The platform information.
    """
    import platform

    r = {}
    for p in dir(platform):
        if p.startswith("_"):
            continue
        try:
            r[p] = getattr(platform, p)()
        except Exception:
            pass

    def all_empty(x):
        return all(all_empty(v) if isinstance(v, (list, tuple)) else v == "" for v in x)

    for k, v in list(r.items()):
        if isinstance(v, (list, tuple)) and all_empty(v):
            del r[k]

    return r


def gpu_info() -> str | list[dict[str, Any]]:
    """Get the GPU information.

    Returns
    -------
    str or list of dict
        The GPU information or an error message.
    """
    import nvsmi

    if not nvsmi.is_nvidia_smi_on_path():
        return "nvdia-smi not found"

    try:
        return [json.loads(gpu.to_json()) for gpu in nvsmi.get_gpus()]
    except subprocess.CalledProcessError as e:
        return e.output.decode("utf-8").strip()


def path_md5(path: str) -> str:
    """Calculate the MD5 hash of a file.

    Parameters
    ----------
    path : str
        The path to the file.

    Returns
    -------
    str
        The MD5 hash of the file.
    """
    import hashlib

    hash = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            hash.update(chunk)
    return hash.hexdigest()


def assets_info(paths: list[str]) -> dict[str, Any]:
    """Get information about the given assets.

    Parameters
    ----------
    paths : list of str
        The list of paths to the assets.

    Returns
    -------
    dict
        The information about the assets.
    """
    result = {}

    for path in paths:
        try:
            mode, ino, dev, nlink, uid, gid, size, atime, mtime, ctime = os.stat(path)  # noqa: F841
            md5 = path_md5(path)
        except Exception as e:
            result[path] = str(e)
            continue

        result[path] = dict(
            size=size,
            atime=datetime.datetime.fromtimestamp(atime).isoformat(),
            mtime=datetime.datetime.fromtimestamp(mtime).isoformat(),
            ctime=datetime.datetime.fromtimestamp(ctime).isoformat(),
            md5=md5,
        )

        try:
            from .checkpoint import peek

            result[path]["peek"] = peek(path)
        except Exception:
            pass

    return result


def gather_provenance_info(assets: list[str] = [], full: bool = False) -> dict[str, Any]:
    """Gather information about the current environment.

    Parameters
    ----------
    assets : list, optional
        A list of file paths for which to collect the MD5 sum, the size and time attributes, by default []
    full : bool, optional
        If true, will also collect various paths, by default False

    Returns
    -------
    dict
        A dictionary with the collected information
    """
    executable = sys.executable

    versions, git_versions = module_versions(full)

    if not full:
        return dict(
            time=datetime.datetime.utcnow().isoformat(),
            python=f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
            module_versions=versions,
            distribution_names=import_name_to_distribution_name(list(versions.keys())),
            git_versions=git_versions,
        )
    else:
        return dict(
            time=datetime.datetime.utcnow().isoformat(),
            executable=executable,
            args=sys.argv,
            python_path=sys.path,
            config_paths=sysconfig.get_paths(),
            module_versions=versions,
            distribution_names=import_name_to_distribution_name(list(versions.keys())),
            git_versions=git_versions,
            platform=platform_info(),
            gpus=gpu_info(),
            assets=assets_info(assets),
        )
