# Copyright 2025 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.
"""Insert additional annotations to RawResult messages."""

from __future__ import annotations

import functools
import re
from collections.abc import Callable

from recipe_engine import recipe_api

from PB.go.chromium.org.luci.buildbucket.proto import common as common_pb
from PB.recipe_engine import result as result_pb


def _get_reasons(*exceptions: Exception) -> list[str]:
    reasons = []

    for exception in exceptions:
        if isinstance(exception, ExceptionGroup):
            for exc in exception.exceptions:
                reasons.extend(_get_reasons(exc))

        elif isinstance(exception, recipe_api.StepFailure):
            reasons.append(exception.reason)

        else:
            reasons.append(str(exception))

    return reasons


def _trim_summary_markdown(summary_markdown: str) -> str:
    summary_markdown = re.sub(r'\n\n\n+', '\n\n', summary_markdown)
    summary_markdown = re.sub(r'^\n+', '', summary_markdown)
    summary_markdown = re.sub(r'\n\n+$', '\n', summary_markdown)
    return summary_markdown


def wrap_run_steps(f: Callable[..., result_pb.RawResult | None]):
    @functools.wraps(f)
    def _wrapper(
        api: recipe_api.RecipeScriptApi,
        *args,
        **kwargs,
    ) -> result_pb.RawResult:
        experiments = [
            x
            for x in api.raw_result.buildbucket.build.input.experiments
            if x.startswith('pigweed.')
        ]

        experiment_lines: list[str] = []
        if experiments:
            experiment_lines.append('\n**Experiments**:')
            for experiment in experiments:
                experiment_lines.append(f'* {experiment}')
            experiment_lines.append('')

        experiment_part = '\n'.join(experiment_lines)

        # If we're not doing bisection and there are no experiments, then
        # there's nothing to add to the result, so directly pass through the
        # return value.
        if not api.raw_result.bisector.enabled and not experiment_part:
            return f(api, *args, **kwargs)

        infra_failure: list[Exception] = []
        step_failure: list[Exception] = []
        other_exception: list[Exception] = []

        try:
            result = f(api, *args, **kwargs)

            if not result:
                result = result_pb.RawResult(
                    summary_markdown='',
                    status=common_pb.SUCCESS,
                )

            summary_markdown = api.raw_result.bisector.summary(
                f'{result.summary_markdown}\n{experiment_part}',
            )

            return result_pb.RawResult(
                summary_markdown=_trim_summary_markdown(summary_markdown),
                status=result.status,
            )

        except* recipe_api.InfraFailure as exc:
            infra_failure.append(exc)

        except* recipe_api.StepFailure as exc:
            step_failure.append(exc)

        except* Exception as exc:
            other_exception.append(exc)

        all_exceptions = [*infra_failure, *step_failure, *other_exception]

        result: list[str] = []
        result.append('\n\n'.join(_get_reasons(*all_exceptions)))
        result.append(experiment_part)

        summary_markdown = api.raw_result.bisector.summary('\n'.join(result))

        status = common_pb.FAILURE
        if other_exception or infra_failure:
            status = common_pb.INFRA_FAILURE
        if 'LUCI CV no longer needs this Tryjob' in summary_markdown:
            status = common_pb.CANCELED

        return result_pb.RawResult(
            summary_markdown=_trim_summary_markdown(summary_markdown),
            status=status,
        )

    return _wrapper


class RawResultApi(recipe_api.RecipeApi):
    @property
    def bisector(self):
        # Don't make all recipes using this module also explicitly depend on
        # the bisector.
        return self.m.bisector

    @property
    def buildbucket(self):
        # Don't make all recipes using this module also explicitly depend on
        # the buildbucket.
        return self.m.buildbucket