# 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.
"""Module for setting up and managing the Pigweed build environment.

This module provides utilities to initialize a consistent build environment
across different platforms. It handles:
- Setting up a Python virtual environment.
- Installing CIPD packages defined in configuration files.
- Setting essential environment variables (e.g., PW_ROOT, PW_PROJECT_ROOT).
- Platform-specific initializations (e.g., macOS SDK).
- Applying toolchain overrides for compilers like Clang.
- Running `pw doctor` to check environment health.

The primary entry point is the `EnvironmentApi.init()` method, which returns
an `Environment` object. This object can then be used as a context manager
to activate the configured environment for subsequent build steps.

Typical usage:

  # In a recipe:
  # Construct an Options proto if needed, or pass None for defaults.
  # from PB.recipe_modules.pigweed.environment import options as env_options
  # opts = env_options.Options(
  #     relative_pigweed_root='pigweed',
  #     config_file='pigweed.json',
  # )
  env = api.environment.init(
      checkout=your_checkout_context,  # From api.checkout
      options=opts,  # Or None
  )
  with env():
      # Run build steps that require the Pigweed environment
      api.step('build something', ['pw', 'build', 'all'])

The `Environment` object also provides access to the configured environment
variables as attributes (e.g., `env.PW_ROOT`).
"""

from __future__ import annotations

import contextlib
import dataclasses
import json
import pprint
import traceback

from recipe_engine import config_types, engine_types, recipe_api

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

from RECIPE_MODULES.pigweed.checkout import api as checkout_api


@dataclasses.dataclass
class Environment:
    api: recipe_api.RecipeApi
    dir: config_types.Path
    checkout: checkout_api.CheckoutContext
    prefixes: dict[str, list[str]] = dataclasses.field(default_factory=dict)
    suffixes: dict[str, list[str]] = dataclasses.field(default_factory=dict)
    env: dict[str, str] = dataclasses.field(default_factory=dict)
    override_gn_args: dict = dataclasses.field(default_factory=dict)

    @contextlib.contextmanager
    def __call__(self):
        with self.api.macos_sdk():
            # 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,
            ):
                yield self

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

        return self.env.get(name)


def path(
    relative_path: config_types.Path,
    checkout: checkout_api.CheckoutContext,
) -> config_types.Path:
    parts = [x for x in relative_path.split('/') if x != '.']
    if parts:
        return checkout.root.joinpath(*parts)
    else:
        return checkout.root  # pragma: no cover


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

    Environment = Environment

    def _init_platform(self, env: Environment) -> None:
        del env  # Unused.
        if self.m.platform.is_mac:
            with self.m.step.nest('setup platform'):
                # Initialize the Mac SDK but then exit. Future calls will
                # execute very quickly.
                with self.m.macos_sdk():
                    pass

    def _init_misc_vars(
        self,
        env: Environment,
        additional_variables: dict[str, str] | None = None,
    ) -> 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['PW_USE_COLOR'] = ''
        env.env['CLICOLOR'] = '0'  # Formerly on https://bixense.com/clicolors/.
        env.env['NO_COLOR'] = '1'  # See https://no-color.org.
        # This should tell ninja to disable colors based on implementation at
        # https://github.com/ninja-build/ninja/blob/master/src/line_printer.cc#L60.
        env.env['CLICOLOR_FORCE'] = '0'
        env.env['GCC_COLORS'] = ''

        env.env.update(additional_variables or {})

        bb_metadata = {}

        build = self.m.buildbucket.build

        if self.m.led.launched_by_led and not self.m.led.led_build:
            env.env['BUILDBUCKET_ID'] = '0'
            bb_metadata['id'] = '0'
        else:
            env.env['BUILDBUCKET_ID'] = str(build.id)
            bb_metadata['id'] = str(build.id)

        env.env['BUILD_NUMBER'] = str(build.number)
        bb_metadata['number'] = str(build.number)

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

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

        bb_metadata['project'] = build.builder.project or 'project'
        bb_metadata['bucket'] = build.builder.bucket or 'bucket'
        bb_metadata['builder'] = build.builder.builder or 'builder'

        env.env['CCACHE_DIR'] = self.m.path.cache_dir / 'ccache'
        env.env['CTCACHE_DIR'] = self.m.path.cache_dir / 'clang_tidy'
        env.env['GOCACHE'] = self.m.path.cache_dir / 'go'
        env.env['PIP_CACHE_DIR'] = self.m.path.cache_dir / 'pip'

        # See b/441999854. Don't allow builds that have github disabled and
        # builds that don't to share Bazel caches. For now, do this by disabling
        # the local cache for builds that disable github.
        if (
            'pigweed.disable_github' in build.input.experiments
            or 'pigweed.disable_rbe' in build.input.experiments
        ):
            env.env['TEST_TMPDIR'] = self.m.path.mkdtemp()
        else:
            env.env['TEST_TMPDIR'] = self.m.path.cache_dir / 'bazel'

        env.env['TRIGGERING_CHANGES_JSON'] = env.checkout.changes_json

        bb_metadata['tags'] = []
        for kv in build.tags:
            bb_metadata['tags'].append((kv.key, kv.value))

        metadata_path = self.m.path.start_dir / 'buildbucket-metadata.json'
        self.m.file.write_json('write bb metadata', metadata_path, bb_metadata)
        env.env['BUILDBUCKET_METADATA_JSON'] = metadata_path

    def _init_pigweed(
        self,
        checkout: checkout_api.CheckoutContext,
        top_presentation: engine_types.StepPresentation,
        use_constraint_file: bool,
        pigweed_root: config_types.Path,
        options: Options,
        env: Environment,
    ) -> None:
        """Run pw_env_setup."""

        json_file = env.dir / 'vars.json'
        shell_file = env.dir / 'setup.sh'
        venv_dir = env.dir / 'venv'

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

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

        cmd: list[str | config_types.Path] = [
            'python3',
            (
                pigweed_root
                / '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',
            env.dir / 'out',
            '--use-existing-cipd',
        ]

        if not options.not_strict:
            cmd.append('--strict')

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

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

        for f in options.additional_cipd_files:
            f = f.replace('$PW_ROOT', str(pigweed_root))
            cmd.extend(('--additional-cipd-file', f))

        cmd.append('--config-file')
        cmd.append(path(options.config_file or 'pigweed.json', checkout))

        with self.m.step.nest('run pw_env_setup') as pres:
            with self.m.defer.context() as defer:
                with env(), self.m.default_timeout():
                    result = defer(self.m.step, 'pw_env_setup', cmd)

                defer(self.m.file.listdir, 'ls', env.dir, recursive=True)

                defer(
                    self.m.save_logs,
                    dirs=(env.dir,),
                    pres=pres,
                    step_passed=result.is_ok(),
                    step_name='environment-setup',
                )

        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 / 'build_overrides' / '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

        with self.m.step.nest('set variables'):
            for var, value in json_data['set'].items():
                with env():
                    if value:
                        env.env[var] = self.m.path.expandvars(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: Environment) -> None:
        """Checks for a toolchain override and applies it."""

        # Using '$fuchsia/checkout' properties to simplify interface with the
        # Fuchsia Toolchain team.
        fuchsia_build_props = self.m.properties.thaw().get(
            '$fuchsia/checkout', {}
        )
        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 / 'override' / 'clang_toolchain'

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

                elif cas_digest := toolchain_props.get('cas_digest'):
                    with self.m.cas.with_instance(
                        'projects/chromium-swarm/instances/default_instance'
                    ):
                        self.m.cas.download(
                            'download',
                            digest=cas_digest,
                            output_dir=toolchain_dir,
                        )

                else:  # pragma: no cover
                    raise KeyError(
                        f'invalid clang toolchain properties: {toolchain_props!r}'
                    )

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

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

    @recipe_api.ignore_warnings('recipe_engine/PYTHON2_DEPRECATED')
    def init(
        self,
        checkout: checkout_api.CheckoutContext,
        options: Options | None = None,
        use_constraint_file: bool = True,
    ) -> Environment:
        pigweed_root = checkout.root
        env = Environment(
            api=self.m,
            dir=checkout.root / 'environment',
            checkout=checkout,
        )

        # 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:
            options = Options()

        if not options.config_file:
            # Always set certain variables even without an environment config.
            self._init_misc_vars(env, options.additional_variables)
            return env

        if options.relative_pigweed_root not in (None, '', '.'):
            pigweed_root = checkout.root / options.relative_pigweed_root

        with self.m.step.nest('environment') as pres:
            with self.m.step.nest('options') as options_pres:
                options_pres.step_summary_text = repr(options)

            try:
                cfg_json = self.m.file.read_json(
                    'read config',
                    path(options.config_file, checkout),
                    test_data={
                        'pw': {
                            'pw_env_setup': {
                                'relative_pigweed_root': 'pigweed',
                            },
                        },
                    },
                )
            except json.decoder.JSONDecodeError as e:  # pragma: no cover
                failure = f'{options.config_file}: {e}'
                pres = self.m.step.empty(failure).presentation
                pres.step_summary_text = traceback.format_exc()
                raise self.m.step.StepFailure(failure)

            pres.logs['config.json'] = pprint.pformat(cfg_json)

            cfg_json = cfg_json.get('pw', cfg_json)
            cfg_json = cfg_json.get('pw_env_setup', cfg_json)
            if 'relative_pigweed_root' in cfg_json:
                pigweed_root = checkout.root / cfg_json['relative_pigweed_root']

            env.env['PW_PROJECT_ROOT'] = checkout.root
            env.env['PW_ROOT'] = pigweed_root

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

            with env():
                try:
                    self.m.step('doctor', ['python', '-m', 'pw_cli', 'doctor'])
                except self.m.step.StepFailure:
                    if not options.allow_doctor_to_fail:
                        raise

        return env