#!/usr/bin/env python

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

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

from . import (format_csv, format_dot, format_json, format_raw, format_table,
               format_tree, log, mapfile, 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-'))


class DocAction(argparse.Action):
    def __call__(self,
                 parser: argparse.ArgumentParser,
                 namespace: argparse.Namespace,
                 values: Union[str, Sequence[Any], None],
                 option_string: Optional[str] = None) -> None:
        from rich.console import Console
        from rich.markdown import Markdown

        console = Console()
        readme_path = Path(os.path.realpath(os.path.dirname(__file__))) / 'docs' / 'readme.md'
        with open(readme_path) as fd:
            readme = fd.read()
            md = Markdown(readme)
            console.print(md)

        parser.exit()


def main() -> None:
    parser = argparse.ArgumentParser(prog='esp_idf_size',
                                     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', 'dot'],
                              default='table',
                              help='Specify output format: table(text), tree, CSV, JSON, raw or DOT.')

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

    deps_group = parser.add_argument_group('Archive dependencies')

    deps_group.add_argument('--archive-dependencies', '--archive-deps',
                            action='store_true',
                            help='Display dependencies or reverse dependencies for all archives.')

    deps_group.add_argument('--dep-symbols', '--dep-syms',
                            action='store_true',
                            help='Include dependency symbols for the --archive-dependencies option.')

    deps_group.add_argument('--dep-reverse', '--dep-rev',
                            action='store_true',
                            help=('Use reverse dependencies for the --archive-dependencies option. '
                                  'This will show the reverse dependencies of archives, instead '
                                  'of archives dependencies.'))

    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('--no-abbrev',
                        action='store_true',
                        help='Do not abbreviate section and file names.')

    parser.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.'))

    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('--use-flash-size',
                        action='store_true',
                        help=('Show the total flash size as defined in the link map file. '
                              'The actual flash size available for the application depends on factors such as the '
                              'partition size for the application and other flash usage, so the total flash size '
                              'in the link map file might not accurately represent the true available size.'))

    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, except when --archive-dependencies '
                              'is used; otherwise, it is ignored.'))

    parser.add_argument('-F', '--filter',
                        metavar='PATTERN',
                        action='append',
                        help=('Use the provided PATTERN to filter archives, object files, or '
                              'symbols in table, CSV or DOT formats. The pattern can include wildcards: '
                              '"*" - matches any sequence of characters, '
                              '"?" - matches any single character, '
                              '"[seq]" - matches any character in the sequence, '
                              '"[!seq]" - matches any character not in the sequence. '
                              'This option can be used multiple times, functioning as a logical OR.'))

    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.'))

    parser.add_argument('--doc',
                        action=DocAction,
                        nargs=0,
                        help=('Display more comprehensive documentation.'))

    ofile = None
    try:
        args = parser.parse_args()

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

        log.set_console(file_stdout=ofile,
                        file_stderr=None,
                        quiet=args.quiet,
                        no_color=args.no_color,
                        force_terminal=args.force_terminal,
                        debug=args.debug)

        if args.no_abbrev and args.unify:
            # --no-abbrev used along with --unify doesn't make sense, because
            # all entries(sections, archives, ...) are using already abbreviated names
            args.no_abbrev = False
            log.warn('The "--no-abbrev" option cannot be used with the "--unify" option and will be ignored.')

        if args.archive_dependencies and args.diff:
            # --archive-dependencies rely on a cross-reference table, so it cannot be used with the --diff option.
            args.diff = None
            log.warn('The "--diff" option cannot be used with the "--archive-dependencies" option and will be ignored.')

        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'))

        map_file = mapfile.MapFile(args.input_file)
        elf = memorymap.get_elf(args.input_file)
        memmap = memorymap.get(args.input_file, load_symbols, args.use_dwarf, map_file, elf)
        if not args.show_unused:
            memorymap.remove_unused(memmap)
        if not args.use_flash_size:
            memorymap.ignore_flash_size(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)
            if not args.use_flash_size:
                memorymap.ignore_flash_size(memmap)
            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, map_file, elf, args)
        elif args.format == 'json2':
            format_json.show(memmap, map_file, elf, args)
        elif args.format == 'raw':
            format_raw.show(memmap, map_file, elf, args)
        elif args.format == 'csv':
            format_csv.show(memmap, map_file, elf, args)
        elif args.format == 'tree':
            format_tree.show(memmap, map_file, elf, args)
        elif args.format == 'dot':
            format_dot.show(memmap, map_file, elf, args)
    except (memorymap.MemMapException, mapfile.MapFileException) as e:
        log.die(str(e))
    except KeyboardInterrupt:
        sys.exit(1)
    finally:
        if ofile:
            ofile.close()
