# Copyright 2022 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.
"""
bloat is a script which generates a size report card for binary files.
"""

import argparse
import json
import logging
import os
from pathlib import Path
import subprocess
import sys
import tempfile
from typing import Iterable

import pw_cli.log

from pw_bloat.bloaty_config import generate_bloaty_config
from pw_bloat.label import DataSourceMap
from pw_bloat.label_output import (
    BloatTableOutput,
    LineCharset,
    RstOutput,
    AsciiCharset,
)

_LOG = logging.getLogger(__name__)

MAX_COL_WIDTH = 50
BINARY_SIZES_EXTENSION = '.binary_sizes.json'


def parse_args() -> argparse.Namespace:
    """Parses the script's arguments."""

    parser = argparse.ArgumentParser('Generate a size report card for binaries')
    parser.add_argument(
        '--gn-arg-path',
        type=str,
        required=True,
        help='File path to json of binaries',
    )
    parser.add_argument(
        '--single-report',
        action="store_true",
        help='Determine if calling single size report',
    )
    parser.add_argument(
        '--json-key-prefix',
        type=str,
        help='Prefix for json keys in size report, default = target name',
        default=None,
    )
    parser.add_argument(
        '--full-json-summary',
        action="store_true",
        help='Include all levels of data sources in json binary report',
    )
    parser.add_argument(
        '--ignore-unused-labels',
        action="store_true",
        help='Do not include labels with size equal to zero in report',
    )

    return parser.parse_args()


def run_bloaty(
    filename: str,
    config: str,
    base_file: str | None = None,
    data_sources: Iterable[str] = (),
    extra_args: Iterable[str] = (),
) -> bytes:
    """Executes a Bloaty size report on some binary file(s).

    Args:
        filename: Path to the binary.
        config: Path to Bloaty config file.
        base_file: Path to a base binary. If provided, a size diff is performed.
        data_sources: List of Bloaty data sources for the report.
        extra_args: Additional command-line arguments to pass to Bloaty.

    Returns:
        Binary output of the Bloaty invocation.

    Raises:
        subprocess.CalledProcessError: The Bloaty invocation failed.
    """

    try:
        # If running from within a Bazel context, find the Bloaty executable
        # within the project's runfiles.
        # pylint: disable=import-outside-toplevel
        from python.runfiles import runfiles  # type: ignore

        r = runfiles.Create()
        bloaty_path = r.Rlocation("bloaty/bloaty")
    except ImportError:
        # Outside of Bazel, use Bloaty from the system path.
        default_bloaty = 'bloaty'
        bloaty_path = os.getenv('BLOATY_PATH', default_bloaty)

    cmd = [
        bloaty_path,
        '-c',
        config,
        '-d',
        ','.join(data_sources),
        '--domain',
        'vm',
        '-n',
        '0',
        filename,
        *extra_args,
    ]

    if base_file is not None:
        cmd.extend(['--', base_file])

    return subprocess.check_output(cmd)


def basic_size_report(
    elf: Path,
    bloaty_config: Path,
    data_sources: Iterable[str] = (),
    extra_args: Iterable[str] = (),
) -> Iterable[str]:
    """Runs a size report on an ELF file using the specified Bloaty config.

    Arguments:
        elf: The ELF binary on which to run.
        bloaty_config: Path to a Bloaty configuration file.
        data_sources: Hierarchical data sources to display.
        extra_args: Additional command line arguments forwarded to bloaty.

    Returns:
        The bloaty TSV output detailing the size report.
    """
    return (
        run_bloaty(
            str(elf.resolve()),
            str(bloaty_config.resolve()),
            data_sources=data_sources,
            extra_args=extra_args,
        )
        .decode('utf-8')
        .splitlines()
    )


class NoMemoryRegions(Exception):
    """Exception raised if an ELF does not define any memory region symbols."""

    def __init__(self, elf: Path):
        super().__init__(f'ELF {elf} does not define memory region symbols')
        self.elf = elf


def memory_regions_size_report(
    elf: Path,
    data_sources: Iterable[str] = (),
    extra_args: Iterable[str] = (),
) -> Iterable[str]:
    """Runs a size report on an ELF file using pw_bloat memory region symbols.

    Arguments:
        elf: The ELF binary on which to run.
        data_sources: Hierarchical data sources to display.
        extra_args: Additional command line arguments forwarded to bloaty.

    Returns:
        The bloaty TSV output detailing the size report.

    Raises:
        NoMemoryRegions: The ELF does not define memory region symbols.
    """
    with tempfile.NamedTemporaryFile() as bloaty_config:
        with open(elf.resolve(), "rb") as infile, open(
            bloaty_config.name, "w"
        ) as outfile:
            result = generate_bloaty_config(
                infile,
                enable_memoryregions=True,
                enable_utilization=False,
                out_file=outfile,
            )

            if not result.has_memoryregions:
                raise NoMemoryRegions(elf)

        return (
            run_bloaty(
                str(elf.resolve()),
                bloaty_config.name,
                data_sources=data_sources,
                extra_args=extra_args,
            )
            .decode('utf-8')
            .splitlines()
        )


def write_file(filename: str, contents: str, out_dir_file: str) -> None:
    path = os.path.join(out_dir_file, filename)
    with open(path, 'w') as output_file:
        output_file.write(contents)
    _LOG.debug('Output written to %s', path)


def create_binary_sizes_json(
    key_prefix: str,
    data_source_map: DataSourceMap,
    full_json: bool,
    ignore_unused_labels: bool,
) -> str:
    """Creates a binary_sizes.json file content from a list of labels.

    Args:
      key_prefix: Prefix for the json keys.
      data_source_map: Hierarchical structure containing size of sources.
      full_json: Report contains all sources, otherwise just top level.
      ignore_unused_labels: Doesn't include labels of size zero in json.

    Returns:
      A string of content to write to binary_sizes.json file.
    """
    json_content = {}
    if full_json:
        *ds_parents, last = data_source_map.get_ds_names()
        for label in data_source_map.labels():
            key = f'{key_prefix}.'
            for ds_parent, label_parent in zip(ds_parents, label.parents):
                key += f'{ds_parent}.{label_parent}.'
            key += f'{last}.{label.name}'
            if label.size != 0 or not ignore_unused_labels:
                json_content[key] = label.size
    else:
        for label in data_source_map.labels(ds_index=0):
            if label.size != 0 or not ignore_unused_labels:
                json_content[f'{key_prefix}.{label.name}'] = label.size
    return json.dumps(json_content, sort_keys=True, indent=2)


def single_target_output(
    target: str,
    bloaty_config: str,
    target_out_file: str,
    out_dir: str,
    data_sources: Iterable[str],
    extra_args: Iterable[str],
    json_key_prefix: str,
    full_json: bool,
    ignore_unused_labels: bool,
) -> int:
    """Generates size report for a single target.

    Args:
      target: The ELF binary on which to run.
      bloaty_config: Path to Bloaty config file.
      target_out_file: Output file name for the generated reports.
      out_dir: Path to write size reports to.
      data_sources: Hierarchical data sources to display.
      extra_args: Additional command-line arguments to pass to Bloaty.
      json_key_prefix: Prefix for the json keys, uses target name by default.
      full_json: Json report contains all hierarchical data source totals.

    Returns:
        Zero on success.

    Raises:
        subprocess.CalledProcessError: The Bloaty invocation failed.
    """

    try:
        single_output = run_bloaty(
            target,
            bloaty_config,
            data_sources=data_sources,
            extra_args=extra_args,
        )

    except subprocess.CalledProcessError:
        _LOG.error('%s: failed to run size report on %s', sys.argv[0], target)
        return 1

    single_tsv = single_output.decode().splitlines()
    single_report = BloatTableOutput(
        DataSourceMap.from_bloaty_tsv(single_tsv), MAX_COL_WIDTH, LineCharset
    )

    data_source_map = DataSourceMap.from_bloaty_tsv(single_tsv)
    rst_single_report = BloatTableOutput(
        data_source_map,
        MAX_COL_WIDTH,
        AsciiCharset,
        True,
    )

    single_report_table = single_report.create_table()

    # Generates contents for summary printed to binary_sizes.json
    binary_json_content = create_binary_sizes_json(
        json_key_prefix, data_source_map, full_json, ignore_unused_labels
    )

    print(single_report_table)
    write_file(target_out_file, rst_single_report.create_table(), out_dir)
    write_file(f'{target_out_file}.txt', single_report_table, out_dir)
    write_file(
        f'{target_out_file}{BINARY_SIZES_EXTENSION}',
        binary_json_content,
        out_dir,
    )

    return 0


def main() -> int:
    """Program entry point."""

    args = parse_args()
    extra_args = ['--tsv']
    data_sources = ['segment_names', 'symbols']
    gn_arg_dict = {}
    json_file = open(args.gn_arg_path)
    gn_arg_dict = json.load(json_file)
    json_key_prefix = args.json_key_prefix

    if args.single_report:
        single_binary_args = gn_arg_dict['binaries'][0]
        if single_binary_args['source_filter']:
            extra_args.extend(
                ['--source-filter', single_binary_args['source_filter']]
            )
        if single_binary_args['data_sources']:
            data_sources = single_binary_args['data_sources']

        # Use target binary name as json key prefix if none given
        if not json_key_prefix:
            json_key_prefix = single_binary_args['target']

        return single_target_output(
            single_binary_args['target'],
            single_binary_args['bloaty_config'],
            gn_arg_dict['target_name'],
            gn_arg_dict['out_dir'],
            data_sources,
            extra_args,
            json_key_prefix,
            args.full_json_summary,
            args.ignore_unused_labels,
        )

    default_data_sources = ['segment_names', 'symbols']

    diff_report = ''
    rst_diff_report = ''
    for curr_diff_binary in gn_arg_dict['binaries']:
        curr_extra_args = extra_args.copy()
        data_sources = default_data_sources

        if curr_diff_binary['source_filter']:
            curr_extra_args.extend(
                ['--source-filter', curr_diff_binary['source_filter']]
            )

        if curr_diff_binary['data_sources']:
            data_sources = curr_diff_binary['data_sources']

        try:
            single_output_base = run_bloaty(
                curr_diff_binary['base'],
                curr_diff_binary['bloaty_config'],
                data_sources=data_sources,
                extra_args=curr_extra_args,
            )

        except subprocess.CalledProcessError:
            _LOG.error(
                '%s: failed to run base size report on %s',
                sys.argv[0],
                curr_diff_binary["base"],
            )
            return 1

        try:
            single_output_target = run_bloaty(
                curr_diff_binary['target'],
                curr_diff_binary['bloaty_config'],
                data_sources=data_sources,
                extra_args=curr_extra_args,
            )

        except subprocess.CalledProcessError:
            _LOG.error(
                '%s: failed to run target size report on %s',
                sys.argv[0],
                curr_diff_binary['target'],
            )
            return 1

        if not single_output_target or not single_output_base:
            continue

        base_dsm = DataSourceMap.from_bloaty_tsv(
            single_output_base.decode().splitlines()
        )
        target_dsm = DataSourceMap.from_bloaty_tsv(
            single_output_target.decode().splitlines()
        )
        diff_dsm = target_dsm.diff(base_dsm)

        diff_report += BloatTableOutput(
            diff_dsm,
            MAX_COL_WIDTH,
            LineCharset,
            diff_label=curr_diff_binary['label'],
        ).create_table()

        curr_rst_report = RstOutput(diff_dsm, curr_diff_binary['label'])
        if rst_diff_report == '':
            rst_diff_report = curr_rst_report.create_table()
        else:
            rst_diff_report += f"{curr_rst_report.add_report_row()}\n"

    print(diff_report)
    write_file(
        gn_arg_dict['target_name'], rst_diff_report, gn_arg_dict['out_dir']
    )
    write_file(
        f"{gn_arg_dict['target_name']}.txt", diff_report, gn_arg_dict['out_dir']
    )

    return 0


if __name__ == '__main__':
    pw_cli.log.install()
    sys.exit(main())
