#!/usr/bin/env python3
# Copyright 2020 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import concurrent.futures
import csv
import json
import logging
import os
import pathlib
import re
import shutil
import subprocess
import sys
from urllib import parse


class StepInfo:
    ABOUT_COLUMNS = (
        "step_name",
        "mem_bytes_inc",
        "mem_bytes_net",
        "top_mem_type",
        "bottom_mem_type",
        "num_objects_inc",
        "num_objects_net",
    )
    # 3 columns separated by " | ".
    _TABLE_ROW_PAT = re.compile(r"([^|]+) \| ([^|]+) \| ([^|]+)$")
    _UNIT_TO_MULTIPLE = {"B": 1, "KB": 1 << 10, "MB": 1 << 20, "GB": 1 << 30}

    def __init__(self, src_path: str):
        self.src_path = src_path.lstrip("/")
        self.step_name = os.path.dirname(src_path.split("/+/u/")[1])
        self.debug_log = None

    def download_log(self):
        self.debug_log = subprocess.check_output(
            ("logdog", "-host", "logs.chromium.org", "cat", self.src_path),
            encoding="utf-8",
        )

    def about(self):
        mem_inc, mem_net = float(0), float(0)
        top_mem_type, bottom_mem_type = ("", 0), ("", 0)
        obj_inc, obj_net = 0, 0
        for line in self.debug_log.split("\n"):
            # Optimization to exit early.
            # This is printed at the end of the memory snapshot.
            if line.endswith("Ends --------"):
                break

            match = self._TABLE_ROW_PAT.match(line)
            if not match:
                continue
            type_name, num_objs, mem = line.split(" | ")
            try:
                num_objs = int(num_objs)
            except ValueError:
                continue

            if num_objs > 0:
                obj_inc += num_objs
            obj_net += num_objs
            mem_magnitude, mem_unit = mem.split()
            num_bytes = float(mem_magnitude) * self._UNIT_TO_MULTIPLE[mem_unit]
            mem_net += num_bytes
            if num_bytes > 0:
                mem_inc += num_bytes

            type_name = type_name.strip()
            if num_bytes > top_mem_type[1]:
                top_mem_type = type_name, num_bytes
            if num_bytes < bottom_mem_type[1]:
                bottom_mem_type = type_name, num_bytes

        return (
            self.step_name,
            str(mem_inc),
            str(mem_net),
            top_mem_type[0],
            bottom_mem_type[0],
            str(obj_inc),
            str(obj_net),
        )


parser = argparse.ArgumentParser()
parser.add_argument("build_url", help="Milo URL for the build of interest")
parser.add_argument("-o", "--output", help="path to which the output will be written")


def main():
    logging.basicConfig(stream=sys.stderr, level=logging.INFO)

    args = parser.parse_args()

    build_url = args.build_url

    if not shutil.which("logdog"):
        sys.exit("logdog not found in $PATH")
    if not shutil.which("bb"):
        sys.exit("bb not found in $PATH")

    build_url_parsed = parse.urlparse(build_url)
    build_id = pathlib.PurePosixPath(build_url_parsed.path).name.lstrip("b")
    logging.info("Getting build %s", build_id)
    build_dict = json.loads(
        subprocess.check_output(
            ("bb", "get", "-json", "-steps", "-p", build_id), encoding="utf-8"
        )
    )
    if (
        not build_dict["input"]["properties"]
        .get("$recipe_engine", {})
        .get("memory_profiler", {})
        .get("enable_snapshot")
    ):
        sys.exit(
            "Build %s does not have memory profiling enabled. See "
            "https://chromium.googlesource.com/infra/luci/recipes-py/+/HEAD/doc/"
            "user_guide.md#memory-leak" % build_url
        )

    step_infos = []
    for step_d in build_dict["steps"]:
        for log_d in step_d.get("logs", ()):
            if log_d["name"] == "$debug":
                step_infos.append(StepInfo(parse.urlparse(log_d["url"]).path))
                break

    logging.info("Downloading debug logs")
    executor = concurrent.futures.ThreadPoolExecutor()
    completed_procs = []
    for step_info in step_infos:
        completed_procs.append(executor.submit(step_info.download_log))
    executor.shutdown(wait=True)
    for cp in completed_procs:
        cp.result()  # Raise any errors.

    logging.info("Writing summary to %s", args.output)
    with open(args.output, "w", encoding="utf-8") as summary_f:
        writer = csv.writer(summary_f)
        writer.writerow(StepInfo.ABOUT_COLUMNS)
        for step_info in step_infos:
            writer.writerow(step_info.about())


if __name__ == "__main__":
    main()