# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

"""
This module specifies the format for search path specifications
in the Pyre configuration.

There are several options that modify how Pyre handles a search
path, see
https://pyre-check.org/docs/configuration/
for details.
"""

import abc
import dataclasses
import glob
import logging
import os
from typing import Dict, Iterable, List, Sequence, Union

from .. import filesystem
from . import exceptions

LOG: logging.Logger = logging.getLogger(__name__)


def _expand_relative_root(path: str, relative_root: str) -> str:
    if not path.startswith("//"):
        return filesystem.expand_relative_path(relative_root, path)
    return path


class Element(abc.ABC):
    @abc.abstractmethod
    def path(self) -> str:
        raise NotImplementedError()

    @abc.abstractmethod
    def command_line_argument(self) -> str:
        raise NotImplementedError()


@dataclasses.dataclass(frozen=True)
class SimpleElement(Element):
    root: str

    def path(self) -> str:
        return self.root

    def command_line_argument(self) -> str:
        return self.root


@dataclasses.dataclass(frozen=True)
class SubdirectoryElement(Element):
    root: str
    subdirectory: str

    def path(self) -> str:
        return os.path.join(self.root, self.subdirectory)

    def command_line_argument(self) -> str:
        return self.root + "$" + self.subdirectory


@dataclasses.dataclass(frozen=True)
class SitePackageElement(Element):
    site_root: str
    package_name: str
    is_toplevel_module: bool = False

    def package_path(self) -> str:
        module_suffix = ".py" if self.is_toplevel_module else ""
        return self.package_name + module_suffix

    def path(self) -> str:
        return os.path.join(self.site_root, self.package_path())

    def command_line_argument(self) -> str:
        return self.site_root + "$" + self.package_path()


class RawElement(abc.ABC):
    @abc.abstractmethod
    def expand_global_root(self, global_root: str) -> "RawElement":
        raise NotImplementedError()

    @abc.abstractmethod
    def expand_relative_root(self, relative_root: str) -> "RawElement":
        raise NotImplementedError()

    @abc.abstractmethod
    def expand_glob(self) -> List["RawElement"]:
        raise NotImplementedError()


@dataclasses.dataclass(frozen=True)
class SimpleRawElement(RawElement):
    root: str

    def expand_global_root(self, global_root: str) -> "SimpleRawElement":
        return SimpleRawElement(
            filesystem.expand_global_root(self.root, global_root=global_root)
        )

    def expand_relative_root(self, relative_root: str) -> "SimpleRawElement":
        return SimpleRawElement(
            _expand_relative_root(self.root, relative_root=relative_root)
        )

    def expand_glob(self) -> List[RawElement]:
        expanded = sorted(glob.glob(self.root))
        if expanded:
            return [SimpleRawElement(path) for path in expanded]
        else:
            LOG.warning(f"'{self.root}' does not match any paths.")
            return []

    def to_element(self) -> SimpleElement:
        return SimpleElement(self.root)


@dataclasses.dataclass(frozen=True)
class SubdirectoryRawElement(RawElement):
    root: str
    subdirectory: str

    def expand_global_root(self, global_root: str) -> "SubdirectoryRawElement":
        return SubdirectoryRawElement(
            root=filesystem.expand_global_root(self.root, global_root=global_root),
            subdirectory=self.subdirectory,
        )

    def expand_relative_root(self, relative_root: str) -> "SubdirectoryRawElement":
        return SubdirectoryRawElement(
            root=_expand_relative_root(self.root, relative_root=relative_root),
            subdirectory=self.subdirectory,
        )

    def expand_glob(self) -> List[RawElement]:
        return [self]

    def to_element(self) -> SubdirectoryElement:
        return SubdirectoryElement(self.root, self.subdirectory)


@dataclasses.dataclass(frozen=True)
class SitePackageRawElement(RawElement):
    package_name: str
    is_toplevel_module: bool = False

    def package_path(self) -> str:
        module_suffix = ".py" if self.is_toplevel_module else ""
        return self.package_name + module_suffix

    def expand_global_root(self, global_root: str) -> "SitePackageRawElement":
        # Site package does not participate in root expansion.
        return self

    def expand_relative_root(self, relative_root: str) -> "SitePackageRawElement":
        # Site package does not participate in root expansion.
        return self

    def expand_glob(self) -> List["RawElement"]:
        return [self]

    def to_element(self, site_root: str) -> SitePackageElement:
        return SitePackageElement(site_root, self.package_name, self.is_toplevel_module)


def create_raw_element(json: Union[str, Dict[str, object]]) -> RawElement:
    if isinstance(json, str):
        return SimpleRawElement(json)
    elif isinstance(json, dict):

        def assert_string_item(input: Dict[str, object], name: str) -> str:
            value = input.get(name, None)
            if not isinstance(value, str):
                raise exceptions.InvalidConfiguration(
                    "Invalid search path element. "
                    f"Expected item `{name}` to be a string but got {value}"
                )
            return value

        if "root" in json and "subdirectory" in json:
            return SubdirectoryRawElement(
                root=assert_string_item(json, "root"),
                subdirectory=assert_string_item(json, "subdirectory"),
            )
        elif "import_root" in json and "source" in json:
            return SubdirectoryRawElement(
                root=assert_string_item(json, "import_root"),
                subdirectory=assert_string_item(json, "source"),
            )
        elif "site-package" in json:
            is_toplevel_module = (
                "is_toplevel_module" in json and json["is_toplevel_module"]
            )
            if not isinstance(is_toplevel_module, bool):
                raise exceptions.InvalidConfiguration(
                    "Invalid search path element. "
                    "Expected `is_toplevel_module` to be a boolean but "
                    f"got {is_toplevel_module}"
                )
            return SitePackageRawElement(
                package_name=assert_string_item(json, "site-package"),
                is_toplevel_module=bool(is_toplevel_module),
            )

    raise exceptions.InvalidConfiguration(
        f"Invalid JSON format for search path element: {json}"
    )


def process_raw_elements(
    raw_elements: Iterable[RawElement],
    site_roots: Sequence[str],
    required: bool = False,
) -> List[Element]:
    elements: List[Element] = []

    def add_if_exists(element: Element) -> bool:
        if os.path.exists(element.path()):
            elements.append(element)
            return True
        return False

    for raw_element in raw_elements:
        expanded_raw_elements = raw_element.expand_glob()
        if len(expanded_raw_elements) == 0 and required:
            raise exceptions.InvalidConfiguration(
                f"Invalid path {raw_element}: does not exist."
            )
        for expanded_raw_element in expanded_raw_elements:
            if isinstance(expanded_raw_element, SitePackageRawElement):
                added = False
                for site_root in site_roots:
                    if added := add_if_exists(
                        expanded_raw_element.to_element(site_root)
                    ):
                        break
                if not added:
                    if required:
                        raise exceptions.InvalidConfiguration(
                            f"Invalid path {expanded_raw_element.package_name}: does not exist."
                        )
                    else:
                        LOG.warning(
                            "Site package does not exist: "
                            f"`{expanded_raw_element.package_name}`"
                        )
            elif isinstance(
                expanded_raw_element, (SimpleRawElement, SubdirectoryRawElement)
            ):
                element = expanded_raw_element.to_element()
                added = add_if_exists(element)
                if not added:
                    if required:
                        raise exceptions.InvalidConfiguration(
                            f"Path does not exist for search path: {element}"
                        )
                    else:
                        LOG.warning(f"Path does not exist for search path: {element}")
            else:
                raise RuntimeError(
                    f"Unhandled raw search path element type: {expanded_raw_element}"
                )

    return elements
