# 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 dataclasses
from typing import Any, Dict, Optional, Sequence

from PB.recipe_modules.pigweed.pw_presubmit import options as options_pb2
from recipe_engine import config_types, recipe_api

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


@dataclasses.dataclass
class Step:
    name: str
    dir: config_types.Path
    substeps: Sequence = dataclasses.field(default_factory=tuple)
    export_dir_name: Optional[str] = dataclasses.field(default=None)
    metadata: Dict = dataclasses.field(default_factory=dict)

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


@dataclasses.dataclass
class PresubmitContext:
    options: options_pb2.Options
    root: config_types.Path
    checkout_root: config_types.Path
    time_rng_seed: int
    _step_objects: Dict = dataclasses.field(
        default_factory=collections.OrderedDict
    )
    list_steps_file: Optional[Any] = dataclasses.field(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, step):
        return Step(
            name=step['name'],
            dir=ctx.root / step['name'],
            substeps=step.get('substeps', ()),
            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(
            options=options or self._options,
            checkout_root=checkout_root,
            root=root or (checkout_root / 'p'),
            time_rng_seed=self.m.time.ms_since_epoch(),
        )

        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') as pres:
            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)

            list_steps_data = 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(
                    {
                        'all_files': ['foo.cc', 'foo.h'],
                        'steps': [{'name': x} for x in test_steps],
                    },
                ),
            ).stdout
            if isinstance(list_steps_data, dict):
                program_steps = list_steps_data['steps']
            else:
                program_steps = list_steps_data

            # 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 / 'list_steps_file.json'
                self.m.file.write_json(
                    'write list steps file',
                    ctx.list_steps_file,
                    list_steps_data,
                )

            # 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),
                )

            pres.step_summary_text = '\n'.join(x['name'] for x in program_steps)

        return ctx

    def _run(
        self, ctx, args, name='run', use_debug_log=True, substep=None, **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')

        if ctx.options.use_time_for_rng_seed:
            cmd.extend(('--rng-seed', ctx.time_rng_seed))

        cmd.extend(args)
        if substep:
            cmd.extend(('--substep', substep))

        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 / 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}')

            for gn_arg in ctx.options.override_gn_arg:
                args.extend(('--override-gn-arg', gn_arg))

            with self.m.step.defer_results():
                if step.substeps:
                    for substep in step.substeps:
                        result = self._run(
                            ctx, args, name=substep, substep=substep
                        )
                        if not result.is_ok:
                            break
                else:
                    result = self._run(ctx, args, name=step.name)
                self.m.gerrit_comment.maybe_post(
                    ctx.options.gerrit_comment, result
                )

                if log_dir:
                    step_log_dir = log_dir / step.name
                else:
                    log_dir = step.export_dir

                if step.export_dir:
                    self.m.file.ensure_directory(
                        f'mkdir {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.save_logs((step.dir,), log_dir, pres=pres)

                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