# Copyright 2020 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.
"""Environment utility functions.

Usage:
env = api.environment.initialize(checkout_root, env_options)
with env():
  ...
"""

import contextlib
import pprint

from PB.recipe_modules.pigweed.environment.options import Options

import attr
from recipe_engine import recipe_api


@attr.s
class Environment:
    _api = attr.ib()
    dir = attr.ib()
    checkout_root = attr.ib()
    prefixes = attr.ib(default=attr.Factory(dict))
    suffixes = attr.ib(default=attr.Factory(dict))
    env = attr.ib(default=attr.Factory(dict))
    override_gn_args = attr.ib(default=attr.Factory(dict))

    @contextlib.contextmanager
    def __call__(self):
        # Using reversed() because things that are added later in environment
        # setup need to override things that came earlier.
        with self._api.context(
            env_prefixes={k: reversed(v) for k, v in self.prefixes.items()},
            env_suffixes=self.suffixes,
            env=self.env,
        ):
            with self._api.macos_sdk():
                yield self

    def __getattr__(self, name):
        if name not in self.env:
            raise AttributeError(name)

        return self.env.get(name)


class EnvironmentApi(recipe_api.RecipeApi):
    """Environment utility functions."""

    def _init_platform(self, env):
        if self.m.platform.is_mac:
            with self.m.step.nest('setup platform'):
                with self.m.macos_sdk():
                    pass

    def _init_misc_vars(self, env, additional_variables=None):
        env.env['PW_ENVIRONMENT_NO_ERROR_ON_UNRECOGNIZED'] = '1'
        env.env['PW_ENVSETUP_DISABLE_SPINNER'] = '1'
        env.env['PW_PRESUBMIT_DISABLE_SUBPROCESS_CAPTURE'] = '1'

        env.env.update(additional_variables or {})

        if self.m.led.launched_by_led:
            # Not using self.m.buildbucket_util.id because some downstream
            # projects need this id to never be longer than a typical
            # buildbucket id. Shorter is fine.
            env.env['BUILDBUCKET_ID'] = '0'
            env.env['BUILD_NUMBER'] = str(self.m.buildbucket.build.number)

        else:
            env.env['BUILDBUCKET_ID'] = str(self.m.buildbucket.build.id)
            env.env['BUILD_NUMBER'] = str(self.m.buildbucket.build.number)

        env.env['BUILDBUCKET_NAME'] = ':'.join(
            (
                self.m.buildbucket.build.builder.project,
                self.m.buildbucket.build.builder.bucket,
                self.m.buildbucket.build.builder.builder,
            )
        )

        if env.env['BUILDBUCKET_NAME'] == '::':
            env.env['BUILDBUCKET_NAME'] = 'project:bucket:builder'

        env.env['CTCACHE_DIR'] = self.m.path['cache'].join('clang_tidy')
        env.env['GOCACHE'] = self.m.path['cache'].join('go')
        env.env['PIP_CACHE_DIR'] = self.m.path['cache'].join('pip')
        env.env['TEST_TMPDIR'] = self.m.path['cache'].join('bazel')

    def _init_pigweed(
        self,
        checkout_root,
        top_presentation,
        use_constraint_file,
        pigweed_root,
        config_file,
        skip_submodule_check,
        env,
    ):
        """Run pw_env_setup."""

        def path(relative_path):
            parts = [
                x for x in relative_path.split('/') if x not in ('.', u'.')
            ]
            if parts:
                return checkout_root.join(*parts)
            else:
                return checkout_root  # pragma: no cover

        json_file = env.dir.join('vars.json')
        shell_file = env.dir.join('setup.sh')
        venv_dir = env.dir.join('venv')

        self.m.file.ensure_directory(
            'mkdir {}'.format(self.m.path.basename(env.dir)), env.dir,
        )

        self.m.file.ensure_directory(
            'mkdir {}'.format(self.m.path.basename(venv_dir)), venv_dir,
        )

        cmd = [
            'python3',
            pigweed_root.join(
                'pw_env_setup', 'py', 'pw_env_setup', 'env_setup.py'
            ),
            '--pw-root',
            pigweed_root,
            '--install-dir',
            env.dir,
            '--json-file',
            json_file,
            '--shell-file',
            shell_file,
            '--virtualenv-gn-out-dir',
            self.m.build.dir,
            '--use-existing-cipd',
            '--strict',
        ]

        if skip_submodule_check:
            cmd.append('--skip-submodule-check')

        if not use_constraint_file:
            cmd.append('--unpin-pip-packages')

        cmd.extend(('--config-file', path(config_file)))

        top_presentation.logs['config.json'] = pprint.pformat(
            self.m.file.read_json('read config', path(config_file))
        )

        with self.m.step.defer_results():
            with self.m.step.nest('run pw_env_setup'):
                with env(), self.m.default_timeout():
                    self.m.step('pw_env_setup', cmd)

                cipd_dir = env.dir.join('cipd')

                log_exts = ('.cfg', '.log', '.sh', '.bat', '.json', '.ensure')

                for directory in (env.dir, cipd_dir, venv_dir):
                    if not self.m.path.isdir(directory):
                        continue

                    with self.m.step.nest(self.m.path.basename(directory)):
                        files = self.m.file.listdir(
                            'ls', directory
                        ).get_result()

                        for entry in files:
                            ext = self.m.path.splitext(entry)[1]
                            if ext not in log_exts:
                                continue

                            self.m.file.read_text(
                                'read {}'.format(self.m.path.basename(entry)),
                                entry,
                            )

        json_data = self.m.file.read_json(
            'read json file',
            json_file,
            test_data={
                'set': {'VIRTUAL_ENV': '/environment/virtualenv'},
                'modify': {
                    'PATH': {'append': ['/environment/bin']},
                    'LD_LIBRARY_PATH': {'prepend': ['/environment/lib']},
                },
            },
        )
        top_presentation.logs['vars.json'] = pprint.pformat(json_data)

        env_gni_path = checkout_root.join('build_overrides').join(
            'pigweed_environment.gni'
        )

        self.m.path.mock_add_file(env_gni_path)
        if self.m.path.isfile(env_gni_path):
            environment_gni = self.m.file.read_text(
                'read gni file', env_gni_path
            )
            top_presentation.logs['pigweed_environment.gni'] = environment_gni

        for var, value in json_data['set'].items():
            env.env[var] = value

        for var, actions in json_data['modify'].items():
            for value in actions.get('prepend', ()):
                env.prefixes.setdefault(var, [])
                env.prefixes[var].append(value)
            for value in actions.get('append', ()):
                env.suffixes.setdefault(var, [])
                env.suffixes[var].append(value)

    def _toolchain_override(self, env):
        """Checks for a toolchain override and applies it."""

        # Using '$fuchsia/build' properties to simplify interface with the
        # Fuchsia Toolchain team.
        fuchsia_build_props = self.m.properties.thaw().get('$fuchsia/build', {})
        toolchain_props = fuchsia_build_props.get('clang_toolchain', {})
        if not toolchain_props:
            return

        with self.m.step.nest('toolchain override'):
            with self.m.context(infra_steps=True):
                toolchain_dir = env.dir.join('override').join('clang_toolchain')

                if toolchain_props['source'] == 'cipd':
                    pkgs = self.m.cipd.EnsureFile()
                    pkgs.add_package(
                        'fuchsia/third_party/clang/${platform}',
                        toolchain_props['version'],
                    )
                    self.m.cipd.ensure(toolchain_dir, pkgs)

                elif toolchain_props['source'] in 'isolated':
                    with self.m.cas.with_instance(
                        'projects/chromium-swarm/instances/default_instance'
                    ):
                        self.m.cas.download(
                            'download',
                            digest=toolchain_props['version'],
                            output_dir=toolchain_dir,
                        )

                else:  # pragma: no cover
                    raise KeyError(
                        f'clang toolchain source {toolchain_props["source"]} '
                        'not recognized'
                    )

                env.prefixes.setdefault('PATH', [])
                env.prefixes['PATH'].append(toolchain_dir)
                env.prefixes['PATH'].append(toolchain_dir.join('bin'))

                clang_prefix = toolchain_dir.join('bin')
                env.override_gn_args[
                    'pw_toolchain_CLANG_PREFIX'
                ] = f'{clang_prefix}/'

    def init(
        self, checkout_root, options=None, use_constraint_file=True,
    ):
        pigweed_root = checkout_root
        env = Environment(
            api=self.m,
            dir=checkout_root.join('environment'),
            checkout_root=checkout_root,
        )

        # If in recipe tests always add at least one variable to make it easier
        # to test use of variables in recipes.
        if self._test_data.enabled:
            env.env['PW_TEST_VAR'] = 'test_value'

        if not options.config_file:
            return env

        if options.relative_pigweed_root not in (None, '', '.'):
            pigweed_root = checkout_root.join(options.relative_pigweed_root)

        if options.root_variable_name:
            env.env[options.root_variable_name] = checkout_root

        with self.m.step.nest('environment') as pres:
            # Setting _initialized immediately because some setup steps need
            # to use the context of previous steps, and invoking env() is
            # the easiest way to do so.
            env.env['PW_ROOT'] = pigweed_root
            env.env['PW_PROJECT_ROOT'] = checkout_root

            self._init_platform(env)
            self._init_misc_vars(env, options.additional_variables)
            self._init_pigweed(
                checkout_root=checkout_root,
                top_presentation=pres,
                use_constraint_file=use_constraint_file,
                pigweed_root=pigweed_root,
                config_file=options.config_file,
                skip_submodule_check=options.skip_submodule_check,
                env=env,
            )
            self._toolchain_override(env)

            with env():
                # If 'pw doctor' fails we can continue, but show the doctor
                # failure in red in the UI.
                try:
                    self.m.step('doctor', ['python', '-m', 'pw_cli', 'doctor'])
                except self.m.step.StepFailure:
                    pass

        return env