# cspell:ignore ftype, chksum
"""Catalog collections within the execution environment."""

from __future__ import annotations

import argparse
import ast
import hashlib
import json
import multiprocessing
import os
import re
import subprocess
import sys

from collections import Counter
from collections import OrderedDict
from datetime import datetime
from datetime import timezone
from json.decoder import JSONDecodeError
from pathlib import Path
from typing import TYPE_CHECKING
from typing import Any

import yaml

from ansible import __version__ as ansible_version
from ansible import plugins
from ansible.utils.plugin_docs import get_docstring
from yaml.error import YAMLError


if TYPE_CHECKING:
    from collections.abc import Iterator


try:
    from yaml import CSafeLoader as SafeLoader
except ImportError:
    from yaml import SafeLoader  # type: ignore

# Import from the source tree whenever possible. When running
# in an execution environment, and therefore not type checking
# import from the key_value_store which was injected in the path.
# The TYPE_CHECKING conditional prevents mypy from attempting the
# import and causing an import error.
try:
    from ansible_navigator.utils.key_value_store import KeyValueStore
except ImportError:
    if not TYPE_CHECKING:
        from key_value_store import KeyValueStore


PROCESSES = (multiprocessing.cpu_count() - 1) or 1


class CollectionCatalog:
    """A collection cataloger."""

    def __init__(self, directories: list[Path]) -> None:
        """Initialize the collection cataloger.

        Args:
            directories: A list of directories that may contain
                collections
        """
        self._directories: list[Path] = directories
        self._collections: OrderedDict[str, dict[Any, Any]] = OrderedDict()
        self._errors: list[dict[str, str]] = []
        self._messages: list[str] = []

    def _catalog_plugins(self, collection: dict[Any, Any]) -> None:
        """Catalog the plugins within a collection.

        Args:
            collection: Details describing the collection
        """
        file_checksums = {}

        file_manifest_file = collection.get("file_manifest_file", {}).get("name")
        if file_manifest_file:
            file_path = Path(collection["path"], file_manifest_file)
            if file_path.exists():
                with file_path.open(encoding="utf-8") as fh:
                    try:
                        loaded = json.load(fh)
                        file_checksums = {v["name"]: v for v in loaded["files"]}
                    except (JSONDecodeError, KeyError) as exc:
                        self._errors.append({"path": str(file_path), "error": str(exc)})

        exempt = ["action", "module_utils", "doc_fragments"]
        plugin_directory = Path(collection["path"], "plugins")
        if plugin_directory.is_dir():
            plugin_dirs = [
                plugin_dir
                for plugin_dir in plugin_directory.iterdir()
                if plugin_dir.is_dir() and plugin_dir.name not in exempt
            ]
            # builtin modules are found in a sibling of the plugin directory
            if collection["known_as"] == "ansible.builtin":
                plugin_dirs.append(Path(collection["path"], "modules"))

            for plugin_dir in plugin_dirs:
                plugin_type = plugin_dir.name
                if plugin_type == "modules":
                    plugin_type = "module"
                files = list(plugin_dir.glob("**/*.py")) + list(plugin_dir.glob("**/*.yml"))
                filenames = iter(files)
                self._process_plugin_dir(
                    plugin_type,
                    filenames,
                    file_checksums,
                    collection,
                )

    def _load_role_argspecs(
        self,
        role: dict[str, str | dict[str, Any]],
        role_directory: Path,
        collection_name: str,
    ) -> None:
        """Load argument specs for a role. Non-fatal on error.

        Args:
            role: The role dict to populate.
            role_directory: Path to the role directory.
            collection_name: The collection name.
        """
        argspec_name = "argument_specs.yml"
        argspec_path = role_directory / "meta" / argspec_name
        role["argument_specs"] = {}
        role["argument_specs_path"] = ""
        error = {"path": str(argspec_path)}
        try:
            with argspec_path.open(encoding="utf-8") as fh:
                parsed = yaml.load(fh, Loader=SafeLoader)
                if not isinstance(parsed, dict) or "argument_specs" not in parsed:
                    error["error"] = f"Malformed {argspec_name} for role in {collection_name}."
                    self._errors.append(error)
                else:
                    role["argument_specs"] = parsed["argument_specs"]
                    role["argument_specs_path"] = str(argspec_path)
        except KeyError:
            error["error"] = f"Malformed {argspec_name} for role in {collection_name}."
            self._errors.append(error)
        except FileNotFoundError:
            error["error"] = f"Failed to find {argspec_name} for role in {collection_name}."
            self._errors.append(error)
        except YAMLError:
            error["error"] = f"Failed to load {argspec_name} for role in {collection_name}."
            self._errors.append(error)

    def _load_role_defaults(
        self,
        role: dict[str, str | dict[str, Any]],
        role_directory: Path,
        collection_name: str,
    ) -> bool:
        """Load defaults for a role.

        Args:
            role: The role dict to populate.
            role_directory: Path to the role directory.
            collection_name: The collection name.

        Returns:
            True if a fatal error occurred.
        """
        defaults_name = "main.yml"
        defaults_path = role_directory / "defaults" / defaults_name
        role["defaults"] = {}
        role["defaults_path"] = ""
        error = {"path": str(defaults_path)}
        try:
            with defaults_path.open(encoding="utf-8") as fh:
                parsed = yaml.load(fh, Loader=SafeLoader)
                role["defaults"] = parsed if isinstance(parsed, dict) else {}
                role["defaults_path"] = str(defaults_path)
        except FileNotFoundError:
            pass
        except YAMLError:
            error["error"] = f"Failed to load {defaults_name} for role in {collection_name}."
            self._errors.append(error)
            return True
        return False

    def _load_role_info(
        self,
        role: dict[str, str | dict[str, Any]],
        role_directory: Path,
        collection_name: str,
    ) -> bool:
        """Load meta/main.yml info for a role.

        Args:
            role: The role dict to populate.
            role_directory: Path to the role directory.
            collection_name: The collection name.

        Returns:
            True if a fatal error occurred.
        """
        meta_name = "main.yml"
        meta_path = role_directory / "meta" / meta_name
        role["info"] = {}
        role["info_path"] = ""
        error = {"path": str(meta_path)}
        try:
            with meta_path.open(encoding="utf-8") as fh:
                parsed = yaml.load(fh, Loader=SafeLoader)
                if not isinstance(parsed, dict):
                    error["error"] = f"Malformed {meta_name} for role in {collection_name}."
                    self._errors.append(error)
                    return True
                role["info"] = parsed
                role["info_path"] = str(meta_path)
        except FileNotFoundError:
            error["error"] = f"Failed to find {meta_name} for role in {collection_name}."
            self._errors.append(error)
            return True
        except YAMLError:
            error["error"] = f"Failed to load {meta_name} for role in {collection_name}."
            self._errors.append(error)
            return True
        return False

    def _load_role_readme(
        self,
        role: dict[str, str | dict[str, Any]],
        role_directory: Path,
        collection_name: str,
    ) -> bool:
        """Load README.md for a role.

        Args:
            role: The role dict to populate.
            role_directory: Path to the role directory.
            collection_name: The collection name.

        Returns:
            True if a fatal error occurred.
        """
        readme_name = "README.md"
        readme_path = role_directory / readme_name
        role["readme"] = ""
        role["readme_path"] = ""
        error = {"path": str(readme_path)}
        try:
            with readme_path.open(encoding="utf-8") as fh:
                role["readme"] = fh.read()
                role["readme_path"] = str(readme_path)
        except FileNotFoundError:
            error["error"] = f"Failed to find {readme_name} for role in {collection_name}."
            self._errors.append(error)
            return True
        return False

    def _catalog_roles(self, collection: dict[str, Any]) -> None:
        """Catalog the roles within a collection.

        Args:
            collection: Details describing the collection
        """
        collection_name: str = collection["known_as"]
        collection["roles"] = []
        roles_directory = Path(collection["path"], "roles")
        if not roles_directory.is_dir():
            return

        for role_directory in roles_directory.iterdir():
            if not role_directory.is_dir():
                continue
            role: dict[str, str | dict[str, Any]] = {
                "short_name": role_directory.name,
                "full_name": f"{collection_name}.{role_directory.name}",
            }
            self._load_role_argspecs(role, role_directory, collection_name)
            error_cataloging_role = self._load_role_defaults(role, role_directory, collection_name)
            error_cataloging_role = (
                self._load_role_info(role, role_directory, collection_name) or error_cataloging_role
            )
            error_cataloging_role = (
                self._load_role_readme(role, role_directory, collection_name)
                or error_cataloging_role
            )
            if not error_cataloging_role:
                collection["roles"].append(role)

    @staticmethod
    def _generate_checksum(file_path: Path, relative_path: Path) -> dict[str, Any]:
        """Generate a standard checksum for a file.

        Args:
            file_path: The path to the file to generate a checksum for
            relative_path: The relative path within the collection
                directory structure

        Returns:
            Details about the file, including the checksum
        """
        sha256_hash = hashlib.sha256()
        with file_path.open("rb") as fh:
            for byte_block in iter(lambda: fh.read(4096), b""):
                sha256_hash.update(byte_block)
        res = {
            "name": relative_path,
            "ftype": "file",
            "chksum_type": "sha256",
            "chksum_sha256": sha256_hash.hexdigest(),
            "format": 1,
        }
        return res

    def _process_plugin_dir(
        self,
        plugin_type: str,
        filenames: Iterator[Path],
        file_checksums: dict[str, dict[Any, Any]],
        collection: dict[Any, Any],
    ) -> None:
        """Process each plugin within one plugin directory.

        Args:
            plugin_type: The type of plugins
            filenames: The filenames of the plugins
            file_checksums: The checksums for the plugin files
            collection: The details of the collection
        """
        for filename in filenames:
            if str(filename.name).startswith("__"):
                continue

            relative_path = Path(filename).relative_to(collection["path"])
            checksum_dict = file_checksums.get(str(relative_path))
            if not checksum_dict:
                checksum_dict = self._generate_checksum(filename, relative_path)
            checksum = checksum_dict[f"chksum_{checksum_dict['chksum_type']}"]
            collection["plugin_checksums"][checksum] = {
                "path": str(relative_path),
                "type": plugin_type,
            }

    def _load_collection_meta(self, directory_path: Path) -> dict[str, Any] | None:
        """Load collection metadata from MANIFEST.json or galaxy.yml.

        Args:
            directory_path: The path to a single collection directory.

        Returns:
            The collection dict, or None if no metadata could be loaded.
        """
        manifest_file = directory_path / "MANIFEST.json"
        galaxy_file = directory_path / "galaxy.yml"
        if manifest_file.exists():
            with manifest_file.open(encoding="utf-8") as fh:
                try:
                    collection = json.load(fh)
                except JSONDecodeError:
                    self._errors.append(
                        {"path": str(manifest_file), "error": "failed to load MANIFEST.json"},
                    )
                else:
                    if not isinstance(collection, dict):
                        self._errors.append(
                            {"path": str(manifest_file), "error": "failed to load MANIFEST.json"},
                        )
                    else:
                        collection["meta_source"] = "MANIFEST.json"
                        return collection
        elif galaxy_file.exists():
            with galaxy_file.open(encoding="utf-8") as fh:
                try:
                    parsed = yaml.load(fh, Loader=SafeLoader)
                except YAMLError:
                    self._errors.append(
                        {"path": str(galaxy_file), "error": "failed to load galaxy.yml"},
                    )
                else:
                    if not isinstance(parsed, dict):
                        self._errors.append(
                            {"path": str(galaxy_file), "error": "failed to load galaxy.yml"},
                        )
                    else:
                        collection = {"collection_info": parsed}
                        collection["meta_source"] = "galaxy.yml"
                        return collection
        return None

    def _finalize_collection(
        self,
        collection: dict[str, Any],
        directory_path: Path,
    ) -> None:
        """Set derived fields and load runtime.yml for a collection.

        Args:
            collection: The collection dict to finalize.
            directory_path: The path to the collection directory.
        """
        collection_name = f"{collection['collection_info']['namespace']}"
        collection_name += f".{collection['collection_info']['name']}"
        collection["known_as"] = collection_name
        collection["plugin_checksums"] = {}
        collection["path"] = str(directory_path)

        runtime_file = directory_path / "meta" / "runtime.yml"
        collection["runtime"] = {}
        if runtime_file.exists():
            with runtime_file.open(encoding="utf-8") as fh:
                try:
                    parsed_runtime = yaml.load(fh, Loader=SafeLoader)
                    collection["runtime"] = (
                        parsed_runtime if isinstance(parsed_runtime, dict) else {}
                    )
                except YAMLError as exc:
                    self._errors.append({"path": str(runtime_file), "error": str(exc)})

    def _one_path(self, directory: Path) -> None:
        """Process the contents of an <...>/ansible_collections/ directory.

        Args:
            directory: The path to collections directory to walk and
                load
        """
        for directory_path in directory.glob("*/*/"):
            collection = self._load_collection_meta(directory_path)
            if collection:
                self._finalize_collection(collection, directory_path)
                self._collections[collection["path"]] = collection
            else:
                msg = (
                    f"collection path '{directory_path}' is ignored as it does not"
                    " have 'MANIFEST.json' and/or 'galaxy.yml' file(s)."
                )
                self._messages.append(msg)

    def _find_shadows(self) -> None:
        """Determine which collections are hidden by another installation of the same."""
        collection_list = list(self._collections.values())
        counts = Counter([collection["known_as"] for collection in collection_list])
        for idx, (collection_path, o_collection) in reversed(
            list(enumerate(self._collections.items())),
        ):
            self._collections[collection_path]["hidden_by"] = []
            if counts[o_collection["known_as"]] > 1:
                for i_collection in reversed(collection_list[0:idx]):
                    if i_collection["known_as"] == o_collection["known_as"]:
                        self._collections[collection_path]["hidden_by"].insert(
                            0,
                            i_collection["path"],
                        )

    def process_directories(self) -> tuple[dict[Any, Any], list[dict[str, str]]]:
        """Process each parent directory.

        Returns:
            All collections found and any errors
        """
        for directory in self._directories:
            collection_directory = directory / "ansible_collections"
            if collection_directory.exists():
                self._one_path(collection_directory)
        self.add_pseudo_builtin()
        for collection in self._collections.values():
            self._catalog_plugins(collection)
            self._catalog_roles(collection)
        self._find_shadows()
        return self._collections, self._errors

    def add_pseudo_builtin(self) -> None:
        """Add the pseudo builtin collection."""
        collection: dict[str, str | list[str] | dict[Any, Any]] = {}
        collection["known_as"] = "ansible.builtin"
        collection["plugin_checksums"] = {}
        collection["path"] = str(Path(plugins.__file__).parents[1])
        collection["runtime"] = {}
        collection["meta_source"] = "None"
        collection["collection_info"] = {"version": ansible_version}
        self._collections[str(collection["path"])] = collection
        msg = f"Added ansible.builtin from: {collection['path']}"
        self._messages.append(msg)


def worker(
    pending_queue: multiprocessing.Queue[Any],
    completed_queue: multiprocessing.Queue[Any],
) -> None:
    """Extract the documentation from a plugin, place in completed queue.

    Args:
        pending_queue: A queue with plugins to process
        completed_queue: The queue in which extracted documentation will
            be placed
    """
    # pylint: disable=import-outside-toplevel
    # pylint: disable=too-many-locals

    # load the fragment_loader _after_ the path is set
    from ansible.plugins.loader import fragment_loader  # noqa: PLC0415

    while True:
        entry = pending_queue.get()
        if entry is None:
            break
        collection_name, checksum, plugin_path = entry

        try:
            if ansible_version.startswith("2.9"):
                (doc, examples, returndocs, metadata) = get_docstring(
                    filename=str(plugin_path),
                    fragment_loader=fragment_loader,
                )
            else:
                (doc, examples, returndocs, metadata) = get_docstring(
                    filename=str(plugin_path),
                    fragment_loader=fragment_loader,
                    collection_name=collection_name,
                )

        except Exception:  # noqa: BLE001
            try:
                with plugin_path.open(mode="r", encoding="utf-8") as f:
                    content = f.read()
                doc, examples, returndocs, metadata = get_doc_withast(content)
                doc, examples, returndocs, metadata = (
                    yaml.load(value, Loader=yaml.SafeLoader)
                    for value in (doc, examples, returndocs, metadata)
                )
            except Exception as exc:  # noqa: BLE001
                err_message = f"{type(exc).__name__} (get_docstring): {exc!s}"
                completed_queue.put(("error", (checksum, plugin_path, err_message)))
                continue

        try:
            q_message = {
                "plugin": {
                    "doc": doc,
                    "examples": examples,
                    "returndocs": returndocs,
                    "metadata": metadata,
                },
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }
            completed_queue.put(("plugin", (checksum, json.dumps(q_message, default=str))))
        except JSONDecodeError as exc:
            err_message = f"{type(exc).__name__} (json_decode_doc): {exc!s}"
            completed_queue.put(("error", (checksum, plugin_path, err_message)))


def identify_missing(
    collections: dict[Any, Any],
    collection_cache: KeyValueStore,
) -> tuple[set[Any], list[Any], int]:
    """Identify plugins missing from the cache.

    Args:
        collections: All plugins found across all collections
        collection_cache: The key value interface to a sqlite database

    Returns:
        Handled and plugins missing from the cache, including a count of
        plugins
    """
    handled = set()
    missing = []
    plugin_count = 0
    for collection in collections.values():
        for checksum, details in collection["plugin_checksums"].items():
            plugin_count += 1
            if checksum not in handled:
                if checksum not in collection_cache:
                    missing.append(
                        (
                            collection["known_as"],
                            checksum,
                            Path(collection["path"], details["path"]),
                        ),
                    )
                handled.add(checksum)
    return handled, missing, plugin_count


def parse_args() -> tuple[argparse.Namespace, list[Path]]:
    # pylint: disable=possibly-used-before-assignment
    """Parse the arguments from the command line.

    Returns:
        The parsed arguments and all directories to search
    """
    parser = argparse.ArgumentParser(description="Catalog collections.")
    parser.add_argument(
        "-d",
        dest="dirs",
        nargs="+",
        help="search within the specified directories",
        default=current_collection_paths,
    )

    parser.add_argument("-a", dest="adjacent", help="prepended to dirs")
    parser.add_argument(
        "-c",
        dest="collection_cache_path",
        help="path to collection cache",
        required=True,
    )
    parsed_args = parser.parse_args()

    adjacent = vars(parsed_args).get("adjacent")
    directories = [adjacent] + parsed_args.dirs if adjacent else parsed_args.dirs

    directories.extend(reversed(sys.path))

    resolved = []
    for directory in directories:
        realpath = Path(directory).resolve()
        if realpath not in resolved:
            resolved.append(realpath)

    return parsed_args, resolved


def retrieve_collections_paths() -> dict[Any, Any]:
    """Retrieve the currently set collection paths.

    Returns:
        Errors or the configured collection directories
    """
    cmd = ["ansible-config", "dump"]
    proc_out = run_command(cmd)
    if "error" in proc_out:
        return proc_out

    # Filter for COLLECTIONS_PATHS instead of using shell pipe and grep
    collections_paths_line = None
    for line in proc_out["stdout"].splitlines():
        if "COLLECTIONS_PATHS" in line:
            collections_paths_line = line
            break

    if not collections_paths_line:
        return {"error": "COLLECTIONS_PATHS not found in ansible-config output"}

    regex = re.compile(r"^(?P<variable>\S+)\((?P<source>.*)\)\s=\s(?P<current>.*)$")
    parsed = regex.match(collections_paths_line)
    if parsed:
        try:
            current = yaml.load(parsed.groupdict()["current"], Loader=SafeLoader)
        except (YAMLError, KeyError) as exc:
            return {"error": str(exc)}
        return {"result": current}
    return {"error": f"corrupt current collection path: {collections_paths_line}"}


def retrieve_docs(
    collection_cache: KeyValueStore,
    errors: list[dict[str, str]],
    missing: list[str],
    stats: dict[Any, Any],
) -> None:
    # pylint: disable=too-many-locals
    """Extract the docs from the plugins.

    Args:
        collection_cache: The key value interface to a sqlite database
        errors: Previous errors encountered
        missing: Plugins missing from the collection cache
        stats: Statistics related to the collection cataloging process
    """
    pending_queue = multiprocessing.Manager().Queue()
    completed_queue = multiprocessing.Manager().Queue()
    processes = []
    for _proc in range(PROCESSES):
        proc = multiprocessing.Process(target=worker, args=(pending_queue, completed_queue))
        processes.append(proc)
        proc.start()

    for entry in missing:
        pending_queue.put(entry)
    for _proc in range(PROCESSES):
        pending_queue.put(None)
    for proc in processes:
        proc.join()

    while not completed_queue.empty():
        message_type, message = completed_queue.get()
        if message_type == "plugin":
            checksum, plugin = message
            collection_cache[checksum] = plugin
            stats["cache_added_success"] += 1
        elif message_type == "error":
            checksum, plugin_path, error = message
            collection_cache[checksum] = json.dumps({"error": error})
            errors.append({"path": str(plugin_path), "error": error})
            stats["cache_added_errors"] += 1


def get_doc_withast(content: Any) -> tuple[Any, Any, Any, Any]:
    """Get the documentation, examples, returndocs, and metadata from the content using ast.

    Args:
        content: The content to parse using ast.

    Returns:
        A tuple containing the documentation, examples, returndocs, and
        metadata.
    """
    doc, examples, returndocs, metadata = "", "", "", ""

    # Parse the content using ast and walk through nodes
    for node in ast.walk(ast.parse(content)):
        if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name):
            target_id = node.targets[0].id

            # Check if node.value is an instance of ast.Constant and its value is a string
            if isinstance(node.value, ast.Constant) and isinstance(node.value.value, str):
                if target_id == "DOCUMENTATION":
                    doc = node.value.value.strip()
                elif target_id == "EXAMPLES":
                    examples = node.value.value.strip()
                elif target_id == "RETURN":
                    returndocs = node.value.value.strip()
                elif target_id == "METADATA":
                    metadata = node.value.value.strip()
    return doc, examples, returndocs, metadata


def run_command(cmd: list[str]) -> dict[str, str]:
    """Run a command using subprocess.

    Args:
        cmd: The command to run, split

    Returns:
        Errors and the stdout from the command run
    """
    try:
        proc_out = subprocess.run(
            cmd,
            capture_output=True,
            check=True,
            text=True,
            shell=False,
        )
    except subprocess.CalledProcessError as exc:
        return {"error": str(exc)}
    except (FileNotFoundError, PermissionError) as exc:
        return {"error": f"{type(exc).__name__}: {exc}"}
    return {"stdout": proc_out.stdout}


def main() -> dict[Any, Any]:
    # pylint: disable=possibly-used-before-assignment
    """Run the collection catalog process.

    Returns:
        The results from the completed collection cataloging process
    """
    stats = {}
    stats["cache_added_success"] = 0
    stats["cache_added_errors"] = 0

    cc_obj = CollectionCatalog(directories=parent_directories)
    collections, errors = cc_obj.process_directories()
    stats["collection_count"] = len(collections)

    collection_cache_path = Path(args.collection_cache_path).resolve().expanduser()
    collection_cache = KeyValueStore(collection_cache_path)

    handled, missing, plugin_count = identify_missing(collections, collection_cache)
    stats["plugin_count"] = plugin_count
    stats["unique plugins"] = len(handled)
    stats["processed"] = len(missing)

    if missing:
        retrieve_docs(collection_cache, errors, missing, stats)

    cached_checksums = collection_cache.keys()
    stats["cache_length"] = len(collection_cache.keys())

    for collection in collections.values():
        for no_doc in set(collection["plugin_checksums"].keys()) - set(cached_checksums):
            del collection["plugin_checksums"][no_doc]

    collection_cache.close()
    return {
        "collections": collections,
        "errors": errors,
        "stats": stats,
        "messages": cc_obj._messages,  # noqa: SLF001
    }


if __name__ == "__main__":
    start_time = datetime.now(timezone.utc)

    collection_paths = retrieve_collections_paths()
    if "error" in collection_paths:
        sys.exit(collection_paths["error"])
    else:
        current_collection_paths = collection_paths["result"]

    args, parent_directories = parse_args()

    COLLECTION_SCAN_PATHS = ":".join(str(parent) for parent in parent_directories)
    os.environ["ANSIBLE_COLLECTIONS_PATH"] = COLLECTION_SCAN_PATHS

    result = main()
    result["stats"]["duration"] = (datetime.now(timezone.utc) - start_time).total_seconds()
    result["collection_scan_paths"] = COLLECTION_SCAN_PATHS
    print(json.dumps(result, default=str))
