#!/usr/bin/env python
#
# This is a tool that works like debug location coverage calculator.
# It parses the llvm-dwarfdump --statistics output by reporting it
# in a more human readable way.
#

from __future__ import print_function
import argparse
import os
import sys
from json import loads
from math import ceil
from collections import OrderedDict
from subprocess import Popen, PIPE

# This special value has been used to mark statistics that overflowed.
TAINT_VALUE = "tainted"

# Initialize the plot.
def init_plot(plt):
    plt.title("Debug Location Statistics", fontweight="bold")
    plt.xlabel("location buckets")
    plt.ylabel("number of variables in the location buckets")
    plt.xticks(rotation=45, fontsize="x-small")
    plt.yticks()


# Finalize the plot.
def finish_plot(plt):
    plt.legend()
    plt.grid(color="grey", which="major", axis="y", linestyle="-", linewidth=0.3)
    plt.savefig("locstats.png")
    print('The plot was saved within "locstats.png".')


# Holds the debug location statistics.
class LocationStats:
    def __init__(
        self,
        file_name,
        variables_total,
        variables_total_locstats,
        variables_with_loc,
        variables_scope_bytes_covered,
        variables_scope_bytes,
        variables_coverage_map,
    ):
        self.file_name = file_name
        self.variables_total = variables_total
        self.variables_total_locstats = variables_total_locstats
        self.variables_with_loc = variables_with_loc
        self.scope_bytes_covered = variables_scope_bytes_covered
        self.scope_bytes = variables_scope_bytes
        self.variables_coverage_map = variables_coverage_map

    # Get the PC ranges coverage.
    def get_pc_coverage(self):
        if self.scope_bytes_covered == TAINT_VALUE or self.scope_bytes == TAINT_VALUE:
            return TAINT_VALUE
        pc_ranges_covered = int(
            ceil(self.scope_bytes_covered * 100.0) / self.scope_bytes
        )
        return pc_ranges_covered

    # Pretty print the debug location buckets.
    def pretty_print(self):
        if self.scope_bytes == 0:
            print("No scope bytes found.")
            return -1

        pc_ranges_covered = self.get_pc_coverage()
        variables_coverage_per_map = {}
        for cov_bucket in coverage_buckets():
            variables_coverage_per_map[cov_bucket] = None
            if (
                self.variables_coverage_map[cov_bucket] == TAINT_VALUE
                or self.variables_total_locstats == TAINT_VALUE
            ):
                variables_coverage_per_map[cov_bucket] = TAINT_VALUE
            else:
                variables_coverage_per_map[cov_bucket] = int(
                    ceil(self.variables_coverage_map[cov_bucket] * 100.0)
                    / self.variables_total_locstats
                )

        print(" =================================================")
        print("            Debug Location Statistics       ")
        print(" =================================================")
        print("     cov%           samples         percentage(~)  ")
        print(" -------------------------------------------------")
        for cov_bucket in coverage_buckets():
            if (
                self.variables_coverage_map[cov_bucket]
                or self.variables_total_locstats == TAINT_VALUE
            ):
                print(
                    "   {0:10}     {1:8}              {2:3}%".format(
                        cov_bucket,
                        self.variables_coverage_map[cov_bucket],
                        variables_coverage_per_map[cov_bucket],
                    )
                )
            else:
                print(
                    "   {0:10}     {1:8d}              {2:3d}%".format(
                        cov_bucket,
                        self.variables_coverage_map[cov_bucket],
                        variables_coverage_per_map[cov_bucket],
                    )
                )
        print(" =================================================")
        print(
            " -the number of debug variables processed: "
            + str(self.variables_total_locstats)
        )
        print(" -PC ranges covered: " + str(pc_ranges_covered) + "%")

        # Only if we are processing all the variables output the total
        # availability.
        if self.variables_total and self.variables_with_loc:
            total_availability = None
            if (
                self.variables_total == TAINT_VALUE
                or self.variables_with_loc == TAINT_VALUE
            ):
                total_availability = TAINT_VALUE
            else:
                total_availability = int(
                    ceil(self.variables_with_loc * 100.0) / self.variables_total
                )
            print(" -------------------------------------------------")
            print(" -total availability: " + str(total_availability) + "%")
        print(" =================================================")

        return 0

    # Draw a plot representing the location buckets.
    def draw_plot(self):
        from matplotlib import pyplot as plt

        buckets = range(len(self.variables_coverage_map))
        plt.figure(figsize=(12, 8))
        init_plot(plt)
        plt.bar(
            buckets,
            self.variables_coverage_map.values(),
            align="center",
            tick_label=self.variables_coverage_map.keys(),
            label="variables of {}".format(self.file_name),
        )

        # Place the text box with the coverage info.
        pc_ranges_covered = self.get_pc_coverage()
        props = dict(boxstyle="round", facecolor="wheat", alpha=0.5)
        plt.text(
            0.02,
            0.90,
            "PC ranges covered: {}%".format(pc_ranges_covered),
            transform=plt.gca().transAxes,
            fontsize=12,
            verticalalignment="top",
            bbox=props,
        )

        finish_plot(plt)

    # Compare the two LocationStats objects and draw a plot showing
    # the difference.
    def draw_location_diff(self, locstats_to_compare):
        from matplotlib import pyplot as plt

        pc_ranges_covered = self.get_pc_coverage()
        pc_ranges_covered_to_compare = locstats_to_compare.get_pc_coverage()

        buckets = range(len(self.variables_coverage_map))
        buckets_to_compare = range(len(locstats_to_compare.variables_coverage_map))

        fig = plt.figure(figsize=(12, 8))
        ax = fig.add_subplot(111)
        init_plot(plt)

        comparison_keys = list(coverage_buckets())
        ax.bar(
            buckets,
            self.variables_coverage_map.values(),
            align="edge",
            width=0.4,
            label="variables of {}".format(self.file_name),
        )
        ax.bar(
            buckets_to_compare,
            locstats_to_compare.variables_coverage_map.values(),
            color="r",
            align="edge",
            width=-0.4,
            label="variables of {}".format(locstats_to_compare.file_name),
        )
        ax.set_xticks(range(len(comparison_keys)))
        ax.set_xticklabels(comparison_keys)

        props = dict(boxstyle="round", facecolor="wheat", alpha=0.5)
        plt.text(
            0.02,
            0.88,
            "{} PC ranges covered: {}%".format(self.file_name, pc_ranges_covered),
            transform=plt.gca().transAxes,
            fontsize=12,
            verticalalignment="top",
            bbox=props,
        )
        plt.text(
            0.02,
            0.83,
            "{} PC ranges covered: {}%".format(
                locstats_to_compare.file_name, pc_ranges_covered_to_compare
            ),
            transform=plt.gca().transAxes,
            fontsize=12,
            verticalalignment="top",
            bbox=props,
        )

        finish_plot(plt)


# Define the location buckets.
def coverage_buckets():
    yield "0%"
    yield "(0%,10%)"
    for start in range(10, 91, 10):
        yield "[{0}%,{1}%)".format(start, start + 10)
    yield "100%"


# Parse the JSON representing the debug statistics, and create a
# LocationStats object.
def parse_locstats(opts, binary):
    # These will be different due to different options enabled.
    variables_total = None
    variables_total_locstats = None
    variables_with_loc = None
    variables_scope_bytes_covered = None
    variables_scope_bytes = None
    variables_scope_bytes_entry_values = None
    variables_coverage_map = OrderedDict()

    # Get the directory of the LLVM tools.
    llvm_dwarfdump_cmd = os.path.join(os.path.dirname(__file__), "llvm-dwarfdump")
    # The statistics llvm-dwarfdump option.
    llvm_dwarfdump_stats_opt = "--statistics"

    # Generate the stats with the llvm-dwarfdump.
    subproc = Popen(
        [llvm_dwarfdump_cmd, llvm_dwarfdump_stats_opt, binary],
        stdin=PIPE,
        stdout=PIPE,
        stderr=PIPE,
        universal_newlines=True,
    )
    cmd_stdout, cmd_stderr = subproc.communicate()

    # TODO: Handle errors that are coming from llvm-dwarfdump.

    # Get the JSON and parse it.
    json_parsed = None

    try:
        json_parsed = loads(cmd_stdout)
    except:
        print("error: No valid llvm-dwarfdump statistics found.")
        sys.exit(1)

    # TODO: Parse the statistics Version from JSON.

    def init_field(name):
        if json_parsed[name] == "overflowed":
            print('warning: "' + name + '" field overflowed.')
            return TAINT_VALUE
        return json_parsed[name]

    if opts.only_variables:
        # Read the JSON only for local variables.
        variables_total_locstats = init_field(
            "#local vars processed by location statistics"
        )
        variables_scope_bytes_covered = init_field(
            "sum_all_local_vars(#bytes in parent scope covered" " by DW_AT_location)"
        )
        variables_scope_bytes = init_field("sum_all_local_vars(#bytes in parent scope)")
        if not opts.ignore_debug_entry_values:
            for cov_bucket in coverage_buckets():
                cov_category = (
                    "#local vars with {} of parent scope covered "
                    "by DW_AT_location".format(cov_bucket)
                )
                variables_coverage_map[cov_bucket] = init_field(cov_category)
        else:
            variables_scope_bytes_entry_values = init_field(
                "sum_all_local_vars(#bytes in parent scope "
                "covered by DW_OP_entry_value)"
            )
            if (
                variables_scope_bytes_covered != TAINT_VALUE
                and variables_scope_bytes_entry_values != TAINT_VALUE
            ):
                variables_scope_bytes_covered = (
                    variables_scope_bytes_covered - variables_scope_bytes_entry_values
                )
            for cov_bucket in coverage_buckets():
                cov_category = (
                    "#local vars - entry values with {} of parent scope "
                    "covered by DW_AT_location".format(cov_bucket)
                )
                variables_coverage_map[cov_bucket] = init_field(cov_category)
    elif opts.only_formal_parameters:
        # Read the JSON only for formal parameters.
        variables_total_locstats = init_field(
            "#params processed by location statistics"
        )
        variables_scope_bytes_covered = init_field(
            "sum_all_params(#bytes in parent scope covered " "by DW_AT_location)"
        )
        variables_scope_bytes = init_field("sum_all_params(#bytes in parent scope)")
        if not opts.ignore_debug_entry_values:
            for cov_bucket in coverage_buckets():
                cov_category = (
                    "#params with {} of parent scope covered "
                    "by DW_AT_location".format(cov_bucket)
                )
                variables_coverage_map[cov_bucket] = init_field(cov_category)
        else:
            variables_scope_bytes_entry_values = init_field(
                "sum_all_params(#bytes in parent scope covered " "by DW_OP_entry_value)"
            )
            if (
                variables_scope_bytes_covered != TAINT_VALUE
                and variables_scope_bytes_entry_values != TAINT_VALUE
            ):
                variables_scope_bytes_covered = (
                    variables_scope_bytes_covered - variables_scope_bytes_entry_values
                )
            for cov_bucket in coverage_buckets():
                cov_category = (
                    "#params - entry values with {} of parent scope covered"
                    " by DW_AT_location".format(cov_bucket)
                )
                variables_coverage_map[cov_bucket] = init_field(cov_category)
    else:
        # Read the JSON for both local variables and formal parameters.
        variables_total = init_field("#source variables")
        variables_with_loc = init_field("#source variables with location")
        variables_total_locstats = init_field(
            "#variables processed by location statistics"
        )
        variables_scope_bytes_covered = init_field(
            "sum_all_variables(#bytes in parent scope covered " "by DW_AT_location)"
        )
        variables_scope_bytes = init_field("sum_all_variables(#bytes in parent scope)")

        if not opts.ignore_debug_entry_values:
            for cov_bucket in coverage_buckets():
                cov_category = (
                    "#variables with {} of parent scope covered "
                    "by DW_AT_location".format(cov_bucket)
                )
                variables_coverage_map[cov_bucket] = init_field(cov_category)
        else:
            variables_scope_bytes_entry_values = init_field(
                "sum_all_variables(#bytes in parent scope covered "
                "by DW_OP_entry_value)"
            )
            if (
                variables_scope_bytes_covered != TAINT_VALUE
                and variables_scope_bytes_entry_values != TAINT_VALUE
            ):
                variables_scope_bytes_covered = (
                    variables_scope_bytes_covered - variables_scope_bytes_entry_values
                )
            for cov_bucket in coverage_buckets():
                cov_category = (
                    "#variables - entry values with {} of parent scope covered "
                    "by DW_AT_location".format(cov_bucket)
                )
                variables_coverage_map[cov_bucket] = init_field(cov_category)

    return LocationStats(
        binary,
        variables_total,
        variables_total_locstats,
        variables_with_loc,
        variables_scope_bytes_covered,
        variables_scope_bytes,
        variables_coverage_map,
    )


# Parse the program arguments.
def parse_program_args(parser):
    parser.add_argument(
        "--only-variables",
        action="store_true",
        default=False,
        help="calculate the location statistics only for local variables",
    )
    parser.add_argument(
        "--only-formal-parameters",
        action="store_true",
        default=False,
        help="calculate the location statistics only for formal parameters",
    )
    parser.add_argument(
        "--ignore-debug-entry-values",
        action="store_true",
        default=False,
        help="ignore the location statistics on locations with " "entry values",
    )
    parser.add_argument(
        "--draw-plot",
        action="store_true",
        default=False,
        help="show histogram of location buckets generated (requires " "matplotlib)",
    )
    parser.add_argument(
        "--compare",
        action="store_true",
        default=False,
        help="compare the debug location coverage on two files provided, "
        "and draw a plot showing the difference  (requires "
        "matplotlib)",
    )
    parser.add_argument("file_names", nargs="+", type=str, help="file to process")

    return parser.parse_args()


# Verify that the program inputs meet the requirements.
def verify_program_inputs(opts):
    if len(sys.argv) < 2:
        print("error: Too few arguments.")
        return False

    if opts.only_variables and opts.only_formal_parameters:
        print("error: Please use just one --only* option.")
        return False

    if not opts.compare and len(opts.file_names) != 1:
        print("error: Please specify only one file to process.")
        return False

    if opts.compare and len(opts.file_names) != 2:
        print("error: Please specify two files to process.")
        return False

    if opts.draw_plot or opts.compare:
        try:
            import matplotlib
        except ImportError:
            print("error: matplotlib not found.")
            return False

    return True


def Main():
    parser = argparse.ArgumentParser()
    opts = parse_program_args(parser)

    if not verify_program_inputs(opts):
        parser.print_help()
        sys.exit(1)

    binary_file = opts.file_names[0]
    locstats = parse_locstats(opts, binary_file)

    if not opts.compare:
        if opts.draw_plot:
            # Draw a histogram representing the location buckets.
            locstats.draw_plot()
        else:
            # Pretty print collected info on the standard output.
            if locstats.pretty_print() == -1:
                sys.exit(0)
    else:
        binary_file_to_compare = opts.file_names[1]
        locstats_to_compare = parse_locstats(opts, binary_file_to_compare)
        # Draw a plot showing the difference in debug location coverage between
        # two files.
        locstats.draw_location_diff(locstats_to_compare)


if __name__ == "__main__":
    Main()
    sys.exit(0)