#!/usr/bin/env python3
"""
@file footprint.py
@brief Analyzes build artefact for memory (FLASH/RAM) footprint

@section License

SPDX-License-Identifier: GPL-2.0-or-later

Copyright (C) 2021-2026 Oryx Embedded SARL. All rights reserved.

This file is part of CycloneBOOT Open.

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 2
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software Foundation,
Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

@author Oryx Embedded SARL (www.oryx-embedded.com)
@version 2.6.2
"""

import argparse
import subprocess
import json
import re
import os

def format_size(n):
    if abs(n) < 1024:
        return f"{n} B"
    return f"{n / 1024:.2f} KB"

def parse_linker_limits(ld_file):
    limits = {'flash': [], 'ram': []}
    try:
        with open(ld_file, 'r', encoding='utf-8', errors='ignore') as f:
            content = f.read()
        pattern = re.compile(r'(\w+)\s*\(.*?\)\s*:\s*ORIGIN\s*=\s*(0x[0-9a-fA-F]+)\s*,\s*LENGTH\s*=\s*([0-9a-fA-FkKmM]+)')

        def to_bytes(val):
            val = val.lower()
            if 'k' in val: return int(val.replace('k', '')) * 1024
            if 'm' in val: return int(val.replace('m', '')) * 1024 * 1024
            return int(val, 0)

        for match in pattern.finditer(content):
            name, origin, length = match.groups()
            start = int(origin, 16)
            end = start + to_bytes(length)
            if 'FLASH' in name.upper(): limits['flash'].append((start, end))
            elif 'RAM' in name.upper(): limits['ram'].append((start, end))
    except Exception as e:
        print(f"Warning: Could not parse LD file: {e}")
    return limits

def matches(obj_path, target_folders):
    path_str = obj_path.replace("\\", "/").lower()
    for folder in target_folders:
        if not folder: continue
        target = folder.replace("\\", "/").rstrip("/").split("/")[-1].lower()
        if target in path_str:
            return True
    return False

def main():
    p = argparse.ArgumentParser(description="Detailed ELF Footprint Analyzer")
    p.add_argument("--map", required=True)
    p.add_argument("--ld", required=True)
    p.add_argument("--elf", required=True)
    p.add_argument("--output", help="Deprecated: use --output-detailed instead")
    p.add_argument("--output-detailed", help="Path to detailed JSON output file (file-level footprint)")
    p.add_argument("--output-summary", help="Path to summary JSON output file (module-level footprint summary)")
    p.add_argument("--folders", nargs="+", required=True)
    p.add_argument("--limit", type=int, default=20, help="Number of top symbols to show")
    p.add_argument("--verbose", action="store_true", help="Show symbol leaderboard in console")
    args = p.parse_args()

    # Support legacy --output flag
    output_detailed = args.output_detailed or args.output
    output_summary = args.output_summary

    limits = parse_linker_limits(args.ld)

    try:
        out = subprocess.check_output(["arm-none-eabi-size", args.elf], text=True).splitlines()
        size_parts = out[1].split()
        true_flash = int(size_parts[0]) + int(size_parts[1])
        true_ram = int(size_parts[1]) + int(size_parts[2])
    except Exception:
        true_flash, true_ram = 0, 0

    stats = {f: {"flash": 0, "ram": 0, "symbols": []} for f in args.folders}

    # Strict tracking to guarantee each map address string is integrated only once
    processed_lines = set()

    # Base structural filters to stop host debug sections from being misidentified as micro runtime RAM
    EXCLUDED_PREFIXES = (
        '.debug',
        '.comment',
        '.ARM.attributes',
        '.group'
    )

    with open(args.map, "r", encoding="utf-8", errors="ignore") as f:
        lines = f.readlines()

    # Regex capturing memory addresses and hex sizing metrics
    hex_pattern = re.compile(r'(0x[0-9a-fA-F]{8,})\s+(0x[0-9a-fA-F]+)')

    for i, line in enumerate(lines):
        if i in processed_lines:
            continue

        match = hex_pattern.search(line)
        if not match: continue

        addr, size = int(match.group(1), 16), int(match.group(2), 16)
        if size == 0: continue

        # Identify Section/Symbol Name
        parts = line.split()
        section_name = parts[0] if not parts[0].startswith("0x") else "unknown"
        if section_name == "unknown" and i > 0:
            prev_parts = lines[i-1].split()
            if prev_parts and not prev_parts[0].startswith("0x"):
                section_name = prev_parts[0]

        # FIX: Catch and remove non-runtime symbols from calculation tracking blocks
        if section_name.startswith(EXCLUDED_PREFIXES):
            continue

        # Identify Object Path
        obj_path = ""
        for p_val in parts:
            if ".obj" in p_val.lower() or ".o" in p_val.lower():
                obj_path = p_val
                break

        # Multi-line Lookahead context parser verification
        if not obj_path and (i + 1) < len(lines):
            next_line = lines[i+1].strip()
            if ".obj" in next_line.lower() or ".o" in next_line.lower():
                obj_path = next_line.split()[0]
                processed_lines.add(i + 1)  # Consume the line to prevent duplicate counting on next loop index

        if not obj_path: continue

        # Iterate target criteria folders to precisely attribute metrics to a single structural domain owner
        for folder in args.folders:
            if matches(obj_path, [folder]):
                is_flash = any(s <= addr < e for s, e in limits['flash'])
                is_ram = any(s <= addr < e for s, e in limits['ram'])

                # If it didn't land in flash bounds and wasn't filtered out as debug data, fall back to RAM assignment
                mem_type = "FLASH" if is_flash else "RAM"

                if mem_type == "FLASH":
                    stats[folder]["flash"] += size
                else:
                    stats[folder]["ram"] += size

                stats[folder]["symbols"].append({
                    "name": section_name,
                    "size": size,
                    "type": mem_type,
                    "obj": os.path.basename(obj_path)
                })
                break  # Break ensures execution flows cleanly out, eliminating folder multi-counting collisions

    # Save Detailed JSON Report
    if output_detailed:
        output_dir = os.path.dirname(output_detailed)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        with open(output_detailed, 'w', encoding='utf-8') as f_out:
            json.dump(stats, f_out, indent=4)

    # Build Summary Structural Map
    summary_data = {
        "modules": {},
        "totals": {
            "claimed_flash_bytes": 0,
            "claimed_ram_bytes": 0,
            "actual_flash_bytes": true_flash,
            "actual_ram_bytes": true_ram
        }
    }

    all_flash_syms = []
    all_ram_syms = []
    total_f, total_r = 0, 0

    # Human-readable console lines
    console_lines = []
    console_lines.append(f"Footprint Summary:\n{'ORYX Module(s)':<25} | {'Flash':<12} | {'RAM':<12}")
    console_lines.append("-" * 55)

    for folder, data in stats.items():
        total_f += data['flash']
        total_r += data['ram']
        display = folder.replace("\\", "/").rstrip("/").split("/")[-1]

        # Populate structured output object
        summary_data["modules"][display] = {
            "folder_path": folder,
            "flash_bytes": data['flash'],
            "ram_bytes": data['ram']
        }

        console_lines.append(f"{display:<25} | {format_size(data['flash']):<12} | {format_size(data['ram']):<12}")

        for sym in data["symbols"]:
            if sym["type"] == "FLASH": all_flash_syms.append(sym)
            else: all_ram_syms.append(sym)

    summary_data["totals"]["claimed_flash_bytes"] = total_f
    summary_data["totals"]["claimed_ram_bytes"] = total_r

    console_lines.append("-" * 55)
    console_lines.append(f"{'TOTAL CLAIMED':<25} | {format_size(total_f):<12} | {format_size(total_r):<12}")
    console_lines.append(f"{'BINARY ACTUAL':<25} | {format_size(true_flash):<12} | {format_size(true_ram):<12}")

    # Print Summary to Console
    for line in console_lines:
        print(line)

    # Save Summary JSON Report
    if output_summary:
        output_dir = os.path.dirname(output_summary)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        with open(output_summary, 'w', encoding='utf-8') as f_out:
            json.dump(summary_data, f_out, indent=4)

    # Verbose Leaderboard
    if args.verbose:
        all_flash_syms.sort(key=lambda x: x['size'], reverse=True)
        print(f"\n--- TOP {args.limit} FLASH SYMBOLS (CLAIMED) ---")
        print(f"{'Symbol/Section':<40} | {'Size':<10} | {'Object'}")
        for s in all_flash_syms[:args.limit]:
            print(f"{s['name'][:38]:<40} | {format_size(s['size']):<10} | {s['obj']}")

        all_ram_syms.sort(key=lambda x: x['size'], reverse=True)
        if all_ram_syms:
            print(f"\n--- TOP {args.limit} RAM SYMBOLS (CLAIMED) ---")
            print(f"{'Symbol/Section':<40} | {'Size':<10} | {'Object'}")
            for s in all_ram_syms[:args.limit]:
                print(f"{s['name'][:38]:<40} | {format_size(s['size']):<10} | {s['obj']}")

    if output_summary:
        print(f"\nSummary JSON saved to: {output_summary}")
    if output_detailed:
        print(f"Detailed JSON saved to: {output_detailed}")

if __name__ == "__main__":
    main()