"""Trie structure for storing artifact paths."""

import fnmatch
import glob
import os

from plaso.engine import logger
from plaso.engine import path_helper


class TrieNode:
    """Represents a node in the Trie.

    Attributes:
      artifacts_names (list[str]): Names of artifacts associated with this node.
      children (dict[str, TrieNode]): Child nodes, keyed by path segment.
      is_root (bool): True if this is the root node.
      path_separator (str): Path separator used in the Trie, default is '/'.
    """

    def __init__(self, path_separator="/", is_root=False):
        """Initializes a trie node object.

        Args:
          path_separator (Optional[str]): the path separator used in paths stored
              in the Trie, default is '/'.
          is_root (Optional[bool]): True if this node is the root node.
        """
        super().__init__()
        self.artifacts_names = []
        self.children = {}
        self.is_root = is_root
        self.path_separator = path_separator


class ArtifactsTrie:
    """Trie structure for storing artifact paths.

    Attributes:
      artifacts_paths (dict[str, list[str]]): Artifact paths for glob expansion,
          keyed by artifact name.
      root (TrieNode): Root node of the Trie.
    """

    def __init__(self):
        """Initializes an artifact trie object."""
        super().__init__()
        self.artifacts_paths = {}
        self.root = TrieNode(is_root=True)

    def AddPath(self, artifact_name, path, path_separator):
        """Adds a path from an artifact definition to the Trie.

        Args:
          artifact_name (str): name of the artifact.
          path (str): path from the artifact definition.
          path_separator (str): path separator.
        """
        logger.debug(f'Adding path: "{path:s}" to artifact: "{artifact_name:s}"')
        path_list = self.artifacts_paths.setdefault(artifact_name, [])
        path_list.append(path)

        current_node = self.root

        # Add a path separator node if this is a new separator.
        if path_separator not in current_node.children:
            current_node.children[path_separator] = TrieNode(
                path_separator=path_separator
            )
        current_node = current_node.children[path_separator]

        # Handle the case when the input path is equal to the path separator.
        if path == path_separator:
            current_node.artifacts_names.append(artifact_name)
            return

        path_segments = path.strip(path_separator).split(path_separator)
        for path_segment in path_segments:
            # Store the path_separator for each node.
            if not hasattr(current_node, "path_separator"):
                current_node.path_separator = path_separator

            if path_segment not in current_node.children:
                current_node.children[path_segment] = TrieNode(path_separator)
            current_node = current_node.children[path_segment]
        current_node.artifacts_names.append(artifact_name)

    def GetMatchingArtifacts(self, path, path_separator):
        """Retrieves the artifact names that match the given path.

        Args:
            path (str): path to match against the Trie.
            path_separator (str): path separator.

        Returns:
          list[str]: artifact names that match the path.
        """
        # Start at the root's child that matches the path_separator.
        if path_separator not in self.root.children:
            return []

        sub_root_node = self.root.children[path_separator]

        # Handle the case when the input path is equal to the path_separator.
        if path == path_separator:
            matching_artifacts = set()
            if sub_root_node.artifacts_names:
                matching_artifacts.update(sub_root_node.artifacts_names)
            return list(matching_artifacts)
        path_segments = path.strip(path_separator).split(path_separator)
        matching_artifacts = set()
        # Update self.artifacts_paths before starting the search.
        self.artifacts_paths = self._GetArtifactsPaths(sub_root_node)

        self._SearchTrie(
            sub_root_node, "", path_segments, path_separator, matching_artifacts
        )
        return list(matching_artifacts)

    def _SearchTrie(
        self, node, current_path, segments, path_separator, matching_artifacts
    ):
        """Searches the trie for paths matching the given path segments.

        Args:
          node (TrieNode): current trie node being traversed.
          current_path (str): path represented by the current node.
          segments (list[str]): remaining path segments to match.
          path_separator (str): path separator.
          matching_artifacts (set[str]): Set to store matching artifact names.
        """
        if node.artifacts_names:
            for artifact_name in node.artifacts_names:
                for artifact_path in self.artifacts_paths.get(artifact_name, []):
                    # Note that the sanitized path is used here to ensure consistency
                    # with the exported path.
                    # TODO: consider moving path sanitation out of engine.
                    if self._ComparePathIfSanitized(
                        current_path, path_separator, artifact_path, node.path_separator
                    ):
                        matching_artifacts.add(artifact_name)
                    elif glob.has_magic(artifact_path):
                        if self._MatchesGlobPattern(
                            artifact_path, current_path, node.path_separator
                        ):
                            matching_artifacts.add(artifact_name)

        if not segments:
            return

        segment = segments[0]
        remaining_segments = segments[1:]

        # Handle glob characters in the current segment.
        for child_segment, child_node in node.children.items():
            # Compare the sanitized version of the path segment stored in
            # the tree to the path segment from to the tool output as it
            # sanitizes path segments before writing data to disk.
            sanitized_child_segment = path_helper.PathHelper.SanitizePathSegments(
                [child_segment]
            ).pop()

            # If the child is an exact match, continue traversal.
            if segment in (child_segment, sanitized_child_segment):
                custom_path = self._CustomPathJoin(
                    path_separator, current_path, child_segment
                )
                self._SearchTrie(
                    child_node,
                    custom_path,
                    remaining_segments,
                    path_separator,
                    matching_artifacts,
                )
            # If the child is a glob, see if it matches.
            elif glob.has_magic(child_segment):
                if self._MatchesGlobPattern(
                    child_segment, segment, child_node.path_separator
                ):
                    custom_path = self._CustomPathJoin(
                        path_separator, current_path, segment
                    )
                    self._SearchTrie(
                        child_node,
                        custom_path,
                        remaining_segments,
                        path_separator,
                        matching_artifacts,
                    )
                    self._SearchTrie(
                        node,
                        custom_path,
                        remaining_segments,
                        path_separator,
                        matching_artifacts,
                    )

    def _ComparePathIfSanitized(
        self, current_path, path_separator, artifact_path, artifact_path_seperator
    ):
        """Compares a current path with an artifact path, handling sanitization.

        This method checks if the current_path matches the artifact_path,
        considering that the artifact_path might have been sanitized.

        Args:
          current_path (str): The current path being checked.
          path_separator (str): Path separator for the current path.
          artifact_path (str): The artifact path to compare against.
          artifact_path_seperator (str): Path separator for the artifact path.

        Returns:
          bool: True if the current path matches the artifact path or its
              sanitized version, False otherwise.
        """
        artifact_path_segments = self._GetNonEmptyPathSegments(
            artifact_path, artifact_path_seperator
        )
        sanitized_path_segments = path_helper.PathHelper.SanitizePathSegments(
            artifact_path_segments
        )
        return self._GetNonEmptyPathSegments(current_path, path_separator) in [
            artifact_path_segments,
            sanitized_path_segments,
        ]

    def _GetNonEmptyPathSegments(self, path, separator):
        """Splits a path into segments and remove non-empty segments.

        Args:
          path (str): The path string to be split.
          separator (str): The path separator.

        Returns:
          list[str]: A list of non-empty path segments.
        """
        return [s for s in path.split(separator) if s]

    def _GetArtifactsPaths(self, node):
        """Retrieves a mapping of artifact names to their paths.

        Args:
          node (TrieNode): current trie node being traversed.

        Returns:
          dict: dictionary mapping artifact names to their paths.
        """
        artifacts_paths = {}
        self._collect_paths(node, "", artifacts_paths)
        return artifacts_paths

    def _collect_paths(self, node, current_path, artifacts):
        """Collects paths from the trie.

        Args:
          node (TrieNode): current node.
          current_path (str): path leading to this node.
          artifacts (dict): dictionary to store artifact paths.
        """
        if node.artifacts_names:
            for artifact_name in node.artifacts_names:
                path_list = artifacts.setdefault(artifact_name, [])
                path_list.append(current_path)

        for segment, child_node in node.children.items():
            # Ensure the path_separator attribute exists.
            if not hasattr(child_node, "path_separator"):
                child_node.path_separator = node.path_separator

            # Construct the next path segment.
            if current_path == child_node.path_separator:
                # Means it is the root folder, i.e. `/`
                next_path = current_path + segment
            else:
                next_path = current_path + child_node.path_separator + segment

            self._collect_paths(child_node, next_path, artifacts)

    def _CustomPathJoin(self, separator, current_path, new_segment):
        """Joins path components using a custom separator, replacing os.sep.

        Args:
          separator (str): The custom separator to use.
          current_path (str): The current path.
          new_segment (str): The new segment to add to it.

        Returns:
          str: The joined path with the custom separator.
        """
        current_path = current_path.replace(separator, os.sep)
        joined_path = os.path.join(current_path, new_segment)
        return joined_path.replace(os.sep, separator)

    def _MatchesGlobPattern(self, glob_pattern, path, path_separator):
        """Checks if a path matches a given glob pattern.

        Args:
          glob_pattern (str): The glob pattern to match against.
          path (str): The path to check.
          path_separator (str): The path separator used in the glob pattern.

        Returns:
          True if the path matches the glob pattern, False otherwise.
        """
        # Normalize paths using the appropriate separators
        glob_pattern = glob_pattern.strip(path_separator).split(path_separator)
        path = path.strip(path_separator).split(path_separator)

        glob_index = 0
        path_index = 0
        while glob_index < len(glob_pattern) and path_index < len(path):
            if glob_pattern[glob_index] == "**":
                # If ** is the last part, it matches everything remaining
                if glob_index == len(glob_pattern) - 1:
                    return True
                # Move to the next part after **
                glob_index += 1
                # Keep advancing in the path until the next part matches
                while path_index < len(path) and not fnmatch.fnmatch(
                    path[path_index], glob_pattern[glob_index]
                ):
                    path_index += 1
            elif not fnmatch.fnmatch(path[path_index], glob_pattern[glob_index]):
                # Mismatch
                return False
            else:
                glob_index += 1
                path_index += 1

        return glob_index == len(glob_pattern) and path_index == len(path)
