# 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.
"""Wrapper for 'pw presubmit' in the project source tree."""

import collections

import attr
from recipe_engine import recipe_api

METADATA = {
    'binary_sizes': (('target', 12345), ('target.budget', 12346)),
    'test_runtimes': (('target', 200), ('target.max', 250)),
}


@attr.s
class Step:
    _api = attr.ib()
    name = attr.ib()
    dir = attr.ib()
    _export_dir_name = attr.ib(default=None)
    metadata = attr.ib(default=attr.Factory(dict))

    @property
    def export_dir(self):
        if not self._export_dir_name:
            return None  # pragma: no cover
        return self.dir.join(self._export_dir_name)


@attr.s
class PresubmitContext:
    _api = attr.ib()
    options = attr.ib()
    root = attr.ib()
    checkout_root = attr.ib()
    _step_objects = attr.ib(default=attr.Factory(collections.OrderedDict))
    list_steps_file = attr.ib(default=None)

    def add_step(self, name, step):
        self._step_objects[name] = step

    @property
    def steps(self):
        return self._step_objects.values()


class PwPresubmitApi(recipe_api.RecipeApi):
    """Calls to checkout code."""

    def _step(self, ctx, name):
        return Step(
            self.m,
            name,
            ctx.root.join(name),
            export_dir_name=ctx.options.export_dir_name,
        )

    def init(self, checkout_root, options=None, root=None):
        options.command_name = options.command_name or 'python -m pw_cli'

        ctx = PresubmitContext(
            api=self.m,
            options=options or self._options,
            checkout_root=checkout_root,
            root=root or checkout_root.join('p'),
        )

        if not ctx.options.step and not ctx.options.program:
            raise self.m.step.StepFailure('no step or program properties')

        with self.m.step.nest('get steps from programs'):
            args = ['--only-list-steps']
            for program in ctx.options.program:
                args.extend(('--program', program))
            for step in ctx.options.step:
                args.extend(('--step', step))

            test_steps = []
            for program in ctx.options.program:
                test_steps.append(f'{program}_0')
                test_steps.append(f'{program}_1')
            test_steps.extend(ctx.options.step)

            program_steps = self._run(
                ctx,
                args,
                name='get steps',
                use_debug_log=False,
                stdout=self.m.json.output(),
                step_test_data=lambda: self.m.json.test_api.output_stream(
                    [{'name': x} for x in test_steps],
                ),
            ).stdout

            # The 'list steps file' is only written if the above run command
            # was successful. Whether 'pw presubmit' supports JSON output for
            # '--only-list-steps' correlates with whether it supports the
            # '--list-steps-file' argument (pwrev/116576).
            if program_steps is not None:
                ctx.list_steps_file = ctx.root.join('list_steps_file.json')
                self.m.file.write_json(
                    'write list steps file', ctx.list_steps_file, program_steps,
                )

            # TODO(b/253021172) Remove this block. It's here until all
            # projects use the new output format with --only-list-steps.
            if program_steps is None:
                raw_steps = (
                    self._run(
                        ctx,
                        args,
                        name=f'get steps text',
                        stdout=self.m.raw_io.output_text(),
                        step_test_data=lambda: self.m.raw_io.test_api.stream_output_text(
                            '\n'.join(test_steps) + '\n'
                        ),
                        use_debug_log=False,
                    )
                    .stdout.strip()
                    .splitlines()
                )
                program_steps = [{'name': x} for x in raw_steps]

            for step in program_steps:
                ctx.add_step(step['name'], self._step(ctx, step['name']))

        return ctx

    def _run(
        self, ctx, args, name='run', use_debug_log=True, **kwargs,
    ):
        if ctx.options.do_not_use_debug_log or not use_debug_log:
            logging_args = ('--loglevel', 'debug')
        else:
            logging_args = (
                '--debug-log',
                self.m.raw_io.output_text(
                    name='debug.log', add_output_log=True,
                ),
            )

        cmd = ctx.options.command_name.split()
        cmd += [
            '--directory',
            ctx.checkout_root,
            *logging_args,
            'presubmit',
            '--output-directory',
            ctx.root,
        ]

        if ctx.list_steps_file:
            cmd += ['--list-steps-file', ctx.list_steps_file]

        if ctx.options.only_on_changed_files:
            args.extend(('--base', 'HEAD~1'))
        elif not ctx.options.do_not_use_full_argument:
            args.append('--full')

        if ctx.options.continue_after_build_error:
            cmd.append('--continue-after-build-error')

        cmd.extend(args)

        with self.m.default_timeout():
            if self.m.resultdb.enabled:
                return self.m.step(
                    name,
                    self.m.resultdb.wrap(
                        cmd,
                        base_variant={
                            'builder': self.m.buildbucket.builder_name,
                            'step': name,
                        },
                        include=True,
                    ),
                    **kwargs,
                )
            else:
                return self.m.step(name, cmd, **kwargs)

    def _process_metadata(self, step):
        if not step.export_dir:
            return  # pragma: no cover

        for name, test_data in METADATA.items():
            step.metadata.setdefault(name, {})

            json_path = step.export_dir.join(f'{name}.json')

            self.m.path.mock_add_file(json_path)
            if self.m.path.isfile(json_path):
                with self.m.step.nest(self.m.path.basename(json_path)):
                    step.metadata[name] = self.m.file.read_json(
                        'read', json_path, test_data=dict(test_data)
                    )

    def run(self, ctx, step, env=None, log_dir=None):
        with self.m.step.nest(step.name) as pres:
            args = ['--step', step.name]

            if env and env.override_gn_args:
                for key, value in env.override_gn_args.items():
                    args.append('--override-gn-arg')
                    if isinstance(value, str):
                        args.append(f'{key}="{value}"')
                    else:
                        args.append(f'{key}={value!r}')

            with self.m.step.defer_results():
                self._run(ctx, args, name=step.name)

                if log_dir:
                    step_log_dir = log_dir.join(step.name)
                else:
                    log_dir = step.export_dir

                if step.export_dir:
                    self.m.file.ensure_directory(
                        'mkdir {}'.format(ctx.options.export_dir_name),
                        step.export_dir,
                    )
                if log_dir and log_dir != step.export_dir:
                    self.m.file.ensure_directory('create log dir', log_dir)
                self.m.build.save_logs(step.dir, log_dir)

                self.m.file.listdir('ls out', step.dir, recursive=True)

            self._process_metadata(step)

    def build_id(self, ctx):
        command = ctx.options.command_name.split()
        command.extend(['--directory', ctx.checkout_root, 'build-id'])
        step_data = self.m.step(
            'get build id',
            command,
            stdout=self.m.raw_io.output_text(),
            step_test_data=lambda: self.m.raw_io.test_api.stream_output_text(
                '123-1234567890'
            ),
            ok_ret='any',
        )

        namespace = None
        if step_data.exc_result.retcode == 0:
            namespace = step_data.stdout.strip()
            if namespace == '0':
                namespace = None

        return namespace