#!/usr/bin/env python

# SPDX-FileCopyrightText: 2023-2024 Espressif Systems (Shanghai) CO LTD
# SPDX-License-Identifier: Apache-2.0

import argparse
import os
import sys
from typing import Any, Optional, Sequence, Union

from . import (format_csv, format_json, format_raw, format_table, format_tree,
               log, memorymap)


class ToggleAction(argparse.Action):
    def __call__(self,
                 parser: argparse.ArgumentParser,
                 namespace: argparse.Namespace,
                 values: Union[str, Sequence[Any], None],
                 option_string: Optional[str] = None) -> None:
        setattr(namespace, self.dest, option_string and not option_string.startswith('--no-'))


def main() -> None:
    parser = argparse.ArgumentParser(prog='esp_idf_size.ng',
                                     description='This tool displays firmware size information for project built by ESP-IDF')

    parser.add_argument('input_file',
                        metavar='MAP_FILE',
                        help='Path to the link map file generated by the ESP-IDF build system.')

    format_group = parser.add_mutually_exclusive_group()
    format_group.add_argument('--format',
                              choices=['table', 'text', 'tree', 'csv', 'json2', 'raw'],
                              default='table',
                              help='Specify output format: table(text), tree, csv, json2 or raw')

    parser.add_argument('--archives',
                        action='store_true',
                        help='Print per-archive sizes.')

    parser.add_argument('--archive-details', '--archive_details',
                        metavar='ARCHIVE_NAME',
                        help='Print detailed symbols per archive')

    parser.add_argument('--files',
                        action='store_true',
                        help='Print per-file sizes.')

    parser.add_argument('--diff',
                        metavar='MAP_FILE',
                        help='Compare sizes with another project.')

    parser.add_argument('--show-unused',
                        action='store_true',
                        help='Show unused memory types and sections.')

    parser.add_argument('--show-unchanged',
                        action='store_true',
                        help='Show unchanged items for --diff operation.')

    parser.add_argument('--lto', '--no-lto',
                        dest='use_dwarf',
                        action=ToggleAction,
                        nargs=0,
                        help=('Enable or disable usage of DWARF debugging information to identify '
                              'archives for symbols without archive. Intended to be used if LTO is enabled. '
                              'If not specified, detect LTO usage from sdkconfig.json, if available.'))

    parser.add_argument('-d', '--debug',
                        action='store_true',
                        help=('Print debug information. Messages are printed to stderr.'))

    parser.add_argument('-o', '--output-file',
                        metavar='OUTPUT_FILE',
                        help=('Print output to the specified file instead of stdout.'))

    parser.add_argument('-s', '--sort',
                        metavar='COLUMN',
                        default=1,
                        help=('Sort table rows based on specified column number, starting from 0. '
                              'Column can be specified also as negative number, where -1 means last column. '
                              'Default is 1 and column 0, containing row description, cannot be used. '
                              'The name of the column can be utilized in place of its numerical identifier. '
                              'Applies only to table and csv formats, otherwise ignored.'))

    parser.add_argument('--sort-diff',
                        action='store_true',
                        help=('Sort entries based on diff value instead of size.'))

    parser.add_argument('--sort-reverse',
                        action='store_false',
                        help=('Sort entries in reversed order. By default descending order is used.'))

    parser.add_argument('-q', '--quiet',
                        action='store_true',
                        help=('Suppress all output.'))

    parser.add_argument('--no-color',
                        action='store_true',
                        help=('Disable ANSI color escape sequences.'))

    parser.add_argument('--force-terminal',
                        action='store_true',
                        default=bool(os.environ.get('ESP_IDF_SIZE_FORCE_TERMINAL')) or None,
                        help=('Enable terminal control codes even if out is not attached to terminal. '
                              'This option is ignored if used along with the "--output-file" option.'))

    # mutually exclusive options
    # --no-abbrev used along with --unify doesn't make sense, because
    # all entries(sections, archives, ...) are using already abbreviated names
    group = parser.add_mutually_exclusive_group()
    group.add_argument('--no-abbrev',
                       action='store_true',
                       help='Do not abbreviate section and file names.')

    group.add_argument('--unify',
                       action='store_true',
                       help=('Use abbreviated names with aggregated size information. '
                             'For example .dram0.bss and .dram1.bss sections will be reported '
                             'under one .bss section. Archives, object files and symbols will be '
                             'aggregated too. This can be useful for the --diff option when '
                             'comparing project built with different esp-idf versions.'))

    ofile = sys.stdout
    try:
        args = parser.parse_args()

        if args.output_file:
            args.force_terminal = False
            ofile = open(args.output_file, 'w')

        log.set_console(ofile, args.quiet, args.no_color, args.force_terminal, args.debug)

        args.abbrev = not args.no_abbrev
        load_symbols = args.archive_details or args.format == 'raw'

        if args.use_dwarf:
            # We need DWARF only for detailed outputs like archives
            args.use_dwarf = any((args.archive_details, args.archives, args.files, args.format == 'raw'))

        memmap = memorymap.get(args.input_file, load_symbols, args.use_dwarf)
        if not args.show_unused:
            memorymap.remove_unused(memmap)
        if args.diff:
            memmap_ref = memorymap.get(args.diff, load_symbols, args.use_dwarf)
            if not args.show_unused:
                memorymap.remove_unused(memmap_ref)
            memmap = memorymap.diff(memmap, memmap_ref)
            if memmap['target'] != memmap['target_diff']:
                log.warn((f'The target of the reference and other project is '
                          f'{memmap["target"]} and {memmap["target_diff"]}, respectively.'))

        if args.unify:
            memorymap.unify(memmap)

        if args.format in ['table', 'text']:
            format_table.show(memmap, args)
        elif args.format == 'json2':
            format_json.show(memmap, args)
        elif args.format == 'raw':
            format_raw.show(memmap, args)
        elif args.format == 'csv':
            format_csv.show(memmap, args)
        elif args.format == 'tree':
            format_tree.show(memmap, args)
    except KeyboardInterrupt:
        sys.exit(1)
    finally:
        if ofile:
            ofile.close()
