# Copyright 2025 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Fetches all atlas-linter-cq bot findings, putting them into CSVs."""

import argparse
import collections
import csv
import dataclasses
import datetime
import json
import logging
from pathlib import Path
import subprocess

from bot_tools import bot_lints


def date_to_proto(date: datetime.date) -> str:
    """Converts the given date to a protojson-amenable timestamp."""
    # Proto accepts timestamps in RFC3339 format, per
    # https://protobuf.dev/reference/php/api-docs/Google/Protobuf/Timestamp.html
    dt = datetime.datetime.combine(date, datetime.time.min, tzinfo=datetime.UTC)
    raw_iso = dt.isoformat(timespec="seconds")
    # This ends with `+00:00`, but proto wants `Z`
    return raw_iso.split("+")[0] + "Z"


def enumerate_bots(
    start_at: datetime.date, stop_at: datetime.date
) -> list[int]:
    """Returns all successful bots created in the given timeframe.

    Args:
        start_at: include bots created on or after this date.
        stop_at: exclude bots created on or after this date.
    """
    predicate = json.dumps(
        {
            "builder": {
                "project": "chromeos",
                "bucket": "cq",
                "builder": "atlas-linters-cq",
            },
            "createTime": {
                "startTime": date_to_proto(start_at),
                "endTime": date_to_proto(stop_at),
            },
            "status": "SUCCESS",
        }
    )

    all_builds = subprocess.run(
        (
            "bb",
            "ls",
            "-id",
            "-predicate",
            predicate,
        ),
        check=True,
        encoding="utf-8",
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
    ).stdout
    return [int(x) for x in all_builds.splitlines()]


def group_findings_by_date(
    info: list[bot_lints.LinterBotInfo],
) -> dict[datetime.date, list[bot_lints.Finding]]:
    """Groups the given linter bots' findings by the date the bot was run on."""
    grouped = collections.defaultdict(list)
    for bot_info in info:
        grouped[bot_info.create_time.date()].extend(bot_info.findings)
    return grouped


def asciify(s: str) -> str:
    """Replaces non-ascii sequences in `s` with an escape sequence.

    This isn't built to be fully, unambiguously round-trippable; it's just meant
    to roughly preserve information from commit messages when they're 99.9%
    ASCII.

    >>> asciify("\uffff")
    '\\uFFFF'
    """
    return s.encode("ascii", "backslashreplace").decode("ascii")


def write_findings(to_file: Path, findings: list[bot_lints.Finding]) -> None:
    """Writes `findings` to `to_file` as a CSV."""
    with to_file.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[x.name for x in dataclasses.fields(bot_lints.Finding)],
            dialect=csv.unix_dialect,
        )
        writer.writeheader()
        # Write the data rows
        for finding in findings:
            write_dict = dataclasses.asdict(finding)
            # non-ascii should be rare, but just to be sure that it's
            # interpreted properly, replace it with an ascii representation
            write_dict["message"] = asciify(write_dict["message"])
            writer.writerow(write_dict)


def write_grouped_findings(
    out_dir: Path, findings: dict[datetime.date, list[bot_lints.Finding]]
) -> None:
    """Writes `findings` to dated subdirectories of `out_dir`.

    `findings` are individually written as CSVs. The output file paths are
    generally of the form `${out_dir}/YYYY/MM/DD/data.csv`.
    """
    out_dir.mkdir(parents=True, exist_ok=True)

    for date, findings_today in findings.items():
        if not findings_today:
            continue
        write_file = (
            out_dir
            / f"{date.year}"
            / f"{str(date.month).zfill(2)}"
            / f"{str(date.day).zfill(2)}"
            / "data.csv"
        )
        write_file.parent.mkdir(parents=True, exist_ok=False)
        write_findings(write_file, findings_today)


def parse_date(x: str) -> datetime.date:
    """Parses a date for argparse.

    Mostly exists so if the user passes an unparsable one, they see the type is
    `parse_date`, not a lambda of some sort.
    """
    return datetime.datetime.strptime(x, "%Y-%m-%d").date()


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--stop-at",
        required=True,
        type=parse_date,
        help="Grab bot creations before this date, in the format YYYY-MM-DD.",
    )
    parser.add_argument(
        "--start-at",
        required=True,
        type=parse_date,
        help="""
        Grab bot creations on or after this date, in the format YYYY-MM-DD.
        """,
    )
    parser.add_argument(
        "--output-dir", required=True, type=Path, help="Dir to write output to."
    )
    parser.add_argument(
        "--debug", action="store_true", help="Enable debug logging."
    )
    return parser.parse_args(argv)


def main(argv: list[str]) -> None:
    opts = parse_args(argv)
    logging.basicConfig(
        format=">> %(asctime)s: %(levelname)s: %(filename)s:%(lineno)d: "
        "%(message)s",
        level=logging.DEBUG if opts.debug else logging.INFO,
    )

    bots_to_inspect = enumerate_bots(opts.start_at, opts.stop_at)
    if not bots_to_inspect:
        logging.info("No new bots found; quit")
        return

    logging.info("Found %d bots to inspect", len(bots_to_inspect))
    bot_infos = [bot_lints.fetch_bot_info(x) for x in bots_to_inspect]
    grouped_findings = group_findings_by_date(bot_infos)
    logging.info(
        "Found %d findings across %d days",
        sum(len(x) for x in grouped_findings.values()),
        len(grouped_findings),
    )
    write_grouped_findings(opts.output_dir, grouped_findings)