# Copyright 2021 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""Script that invokes protoc to generate code for .proto files."""

import argparse
from collections import defaultdict
from dataclasses import dataclass
from itertools import chain
import json
from pathlib import Path
import sys
import textwrap
from typing import Any, Dict, Iterable, Iterator, List, Set, Sequence, TextIO

try:
    from pw_build.mirror_tree import mirror_paths
except ImportError:
    # Append this path to the module search path to allow running this module
    # before the pw_build package is installed.
    sys.path.append(str(Path(__file__).resolve().parent.parent))
    from pw_build.mirror_tree import mirror_paths


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument('--label', help='Label for this Python package')
    parser.add_argument('--proto-library',
                        dest='proto_libraries',
                        type=argparse.FileType('r'),
                        default=[],
                        action='append',
                        help='Paths')
    parser.add_argument('--generated-root',
                        required=True,
                        type=Path,
                        help='The base directory for the Python package')
    parser.add_argument('--setup-json',
                        required=True,
                        type=argparse.FileType('r'),
                        help='setup.py keywords as JSON')
    parser.add_argument('--module-as-package',
                        action='store_true',
                        help='Generate an __init__.py that imports everything')
    parser.add_argument('files',
                        type=Path,
                        nargs='+',
                        help='Relative paths to the files in the package')
    return parser.parse_args()


def _check_nested_protos(label: str, proto_info: Dict[str, Any]) -> None:
    """Checks that the proto library refers to this package."""
    python_package = proto_info['nested_in_python_package']

    if python_package != label:
        raise ValueError(
            f"{label}'s 'proto_library' is set to {proto_info['label']}, but "
            f"that target's 'python_package' is {python_package or 'not set'}. "
            f"Set {proto_info['label']}'s 'python_package' to {label}.")


@dataclass(frozen=True)
class _ProtoInfo:
    root: Path
    sources: Sequence[Path]
    deps: Sequence[str]


def _collect_all_files(
        root: Path, files: List[Path],
        paths_to_collect: Iterable[_ProtoInfo]) -> Dict[str, Set[str]]:
    """Collects files in output dir, adds to files; returns package_data."""
    root.mkdir(exist_ok=True)

    for proto_info in paths_to_collect:
        # Mirror the proto files to this package.
        files += mirror_paths(proto_info.root, proto_info.sources, root)

    # Find all subpackages, including empty ones.
    subpackages: Set[Path] = set()
    for file in (f.relative_to(root) for f in files):
        subpackages.update(root / path for path in file.parents)
    subpackages.remove(root)

    # Make sure there are __init__.py and py.typed files for each subpackage.
    for pkg in subpackages:
        for file in (pkg / name for name in ['__init__.py', 'py.typed']):
            if not file.exists():
                file.touch()
            files.append(file)

    pkg_data: Dict[str, Set[str]] = defaultdict(set)

    # Add all non-source files to package data.
    for file in (f for f in files if f.suffix != '.py'):
        pkg = file.parent

        package_name = pkg.relative_to(root).as_posix().replace('/', '.')
        pkg_data[package_name].add(file.name)

    return pkg_data


_SETUP_PY_FILE = '''\
# Generated file. Do not modify.
# pylint: skip-file

import setuptools  # type: ignore

setuptools.setup(
{keywords}
)
'''


def _generate_setup_py(pkg_data: dict, keywords: dict) -> str:
    setup_keywords = dict(
        packages=list(pkg_data),
        package_data={pkg: list(files)
                      for pkg, files in pkg_data.items()},
    )

    assert not any(kw in keywords for kw in setup_keywords), (
        'Generated packages may not specify "packages" or "package_data"')
    setup_keywords.update(keywords)

    return _SETUP_PY_FILE.format(keywords='\n'.join(
        f'    {k}={v!r},' for k, v in setup_keywords.items()))


def _import_module_in_package_init(all_files: List[Path]) -> None:
    """Generates an __init__.py that imports the module.

    This makes an individual module usable as a package. This is used for proto
    modules.
    """
    sources = [
        f for f in all_files if f.suffix == '.py' and f.name != '__init__.py'
    ]
    assert len(sources) == 1, (
        'Module as package expects a single .py source file')

    source, = sources
    source.parent.joinpath('__init__.py').write_text(
        f'from {source.stem}.{source.stem} import *\n')


def _load_metadata(label: str,
                   proto_libraries: Iterable[TextIO]) -> Iterator[_ProtoInfo]:
    for proto_library_file in proto_libraries:
        info = json.load(proto_library_file)
        _check_nested_protos(label, info)

        deps = []
        for dep in info['dependencies']:
            with open(dep) as file:
                deps.append(json.load(file)['package'])

        yield _ProtoInfo(Path(info['root']),
                         tuple(Path(p) for p in info['protoc_outputs']), deps)


def main(generated_root: Path, files: List[Path], module_as_package: bool,
         setup_json: TextIO, label: str,
         proto_libraries: Iterable[TextIO]) -> int:
    """Generates a setup.py and other files for a Python package."""
    proto_infos = list(_load_metadata(label, proto_libraries))
    try:
        pkg_data = _collect_all_files(generated_root, files, proto_infos)
    except ValueError as error:
        msg = '\n'.join(textwrap.wrap(str(error), 78))
        print(
            f'ERROR: Failed to generate Python package {label}:\n\n'
            f'{textwrap.indent(msg, "  ")}\n',
            file=sys.stderr)
        return 1

    with setup_json:
        setup_keywords = json.load(setup_json)

    install_requires = setup_keywords.setdefault('install_requires', [])
    install_requires += chain.from_iterable(i.deps for i in proto_infos)

    if module_as_package:
        _import_module_in_package_init(files)

    # Create the setup.py file for this package.
    generated_root.joinpath('setup.py').write_text(
        _generate_setup_py(pkg_data, setup_keywords))

    return 0


if __name__ == '__main__':
    sys.exit(main(**vars(_parse_args())))