#!/usr/bin/env bash
set -euo pipefail

usage() {
    cat <<'EOF'
Usage:
  memory_report.sh --elf <path/to/file.elf> --mcu <mcu> [--su <file.su> ...]
EOF
}

elf_file=""
mcu=""
su_files=()

while (($# > 0)); do
    case "$1" in
        --elf)
            elf_file="${2:-}"
            shift 2
            ;;
        --mcu)
            mcu="${2:-}"
            shift 2
            ;;
        --su)
            su_files+=("${2:-}")
            shift 2
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown argument: $1" >&2
            usage >&2
            exit 1
            ;;
    esac
done

if [[ -z "$elf_file" || -z "$mcu" ]]; then
    usage >&2
    exit 1
fi

if [[ ! -f "$elf_file" ]]; then
    echo "ELF file not found: $elf_file" >&2
    exit 1
fi

for tool in avr-size avr-gcc avr-objdump awk; do
    if ! command -v "$tool" >/dev/null 2>&1; then
        echo "Required tool not found: $tool" >&2
        exit 1
    fi
done

section_size() {
    local section="$1"
    avr-size -A "$elf_file" | awk -v section="$section" '$1 == section { print $2; found = 1 } END { if (!found) print 0 }'
}

estimate_stack() {
    local elf="$1"
    shift
    local su_inputs=("$@")
    local su_existing=()
    local tmp_su

    for su in "${su_inputs[@]}"; do
        if [[ -f "$su" ]]; then
            su_existing+=("$su")
        fi
    done

    tmp_su="$(mktemp)"
    if ((${#su_existing[@]} > 0)); then
        cat "${su_existing[@]}" >"$tmp_su"
    fi

    avr-objdump -d "$elf" | awk -v su_file="$tmp_su" '
function valid_symbol(name) {
    return (name ~ /^[A-Za-z_][A-Za-z0-9_$.@]*$/) && (name !~ /^L0/);
}
function ignore_sp_write(name) {
    return (name == "__ctors_end" || name == "__do_copy_data" || name == "__do_clear_bss");
}
function add_edge(src, dst, i) {
    for (i = 1; i <= edge_count[src]; ++i) {
        if (edge[src, i] == dst) {
            return;
        }
    }
    edge[src, ++edge_count[src]] = dst;
}
function mark_reachable(node, i, child) {
    if (reach_state[node]) {
        return;
    }
    reach_state[node] = 1;
    reachable[node] = 1;
    for (i = 1; i <= edge_count[node]; ++i) {
        child = edge[node, i];
        mark_reachable(child);
    }
}
function compute(node, i, child, best_total, best_path, unk) {
    if (visit_state[node] == 2) {
        return;
    }
    if (visit_state[node] == 1) {
        cycle_node[node] = 1;
        total[node] = self_stack[node];
        path[node] = node " [cycle]";
        unknown[node] = 1;
        visit_state[node] = 2;
        return;
    }

    visit_state[node] = 1;
    best_total = 0;
    best_path = "";
    unk = ((!known[node]) || (node in indirect_call));

    for (i = 1; i <= edge_count[node]; ++i) {
        child = edge[node, i];
        compute(child);
        if (unknown[child]) {
            unk = 1;
        }
        if (total[child] > best_total) {
            best_total = total[child];
            best_path = path[child];
        }
    }

    total[node] = self_stack[node] + best_total;
    if (best_path != "") {
        path[node] = node " -> " best_path;
    } else {
        path[node] = node;
    }
    unknown[node] = unk;
    visit_state[node] = 2;
}
BEGIN {
    while ((getline line < su_file) > 0) {
        split(line, fields, "\t");
        if (length(fields) < 2) {
            continue;
        }
        split(fields[1], loc_parts, ":");
        fn = loc_parts[length(loc_parts)];
        usage = fields[2] + 0;
        qualifier = (length(fields) >= 3 ? fields[3] : "");

        if (!(fn in su_stack) || usage > su_stack[fn]) {
            su_stack[fn] = usage;
        }
        if (qualifier ~ /dynamic/ && qualifier !~ /bounded/) {
            su_unbounded[fn] = 1;
        }
    }
    close(su_file);
}
{
    if (match($0, /^[[:xdigit:]]+ <([^>]+)>:$/, header_match)) {
        sym = header_match[1];
        if (valid_symbol(sym) && sym !~ /^\./) {
            current_fn = sym;
            functions[sym] = 1;
            if (!(sym in disasm_depth)) {
                disasm_depth[sym] = 0;
                disasm_max[sym] = 0;
                disasm_complex[sym] = 0;
            }
        }
    }

    if (current_fn == "") {
        next;
    }

    if ($0 ~ /[[:space:]](call|rcall)[[:space:]]/ && match($0, / <([^>]+)>/, call_match)) {
        callee = call_match[1];
        if (valid_symbol(callee) && callee !~ /^\./) {
            add_edge(current_fn, callee);
            functions[callee] = 1;
        }
    }

    if ($0 ~ /[[:space:]](icall|eicall)[[:space:]]/) {
        indirect_call[current_fn] = 1;
    }

    if ($0 ~ /[[:space:]]push[[:space:]]/) {
        ++disasm_depth[current_fn];
        if (disasm_depth[current_fn] > disasm_max[current_fn]) {
            disasm_max[current_fn] = disasm_depth[current_fn];
        }
    }

    if ($0 ~ /[[:space:]]pop[[:space:]]/) {
        if (disasm_depth[current_fn] > 0) {
            --disasm_depth[current_fn];
        }
    }

    if ($0 ~ /rcall[[:space:]]+\.\+0/) {
        disasm_depth[current_fn] += 2;
        if (disasm_depth[current_fn] > disasm_max[current_fn]) {
            disasm_max[current_fn] = disasm_depth[current_fn];
        }
    }

    if (($0 ~ /out[[:space:]]+0x3[dD]/ || $0 ~ /out[[:space:]]+0x3[eE]/) && !ignore_sp_write(current_fn)) {
        disasm_complex[current_fn] = 1;
    }
}
END {
    for (fn in su_stack) {
        functions[fn] = 1;
    }

    for (fn in functions) {
        if (fn in su_stack) {
            self_stack[fn] = su_stack[fn];
            known[fn] = ((fn in su_unbounded) ? 0 : 1);
        } else if ((fn in disasm_max) && (disasm_complex[fn] == 0)) {
            self_stack[fn] = disasm_max[fn] + 2;
            known[fn] = 1;
        } else {
            self_stack[fn] = 0;
            known[fn] = 0;
        }
    }

    if (!("main" in functions)) {
        print "STACK_ERROR\tmain-not-found";
        exit;
    }

    compute("main");
    main_total = total["main"];
    main_unknown = unknown["main"];
    main_path = path["main"];

    max_isr_total = 0;
    max_isr_name = "";
    max_isr_path = "";
    max_isr_unknown = 0;
    isr_count = 0;

    for (fn in functions) {
        if (fn ~ /^__vector_[0-9]+$/) {
            ++isr_count;
            compute(fn);
            if (total[fn] > max_isr_total) {
                max_isr_total = total[fn];
                max_isr_name = fn;
                max_isr_path = path[fn];
                max_isr_unknown = unknown[fn];
            }
        }
    }

    combined_total = main_total + max_isr_total;
    combined_unknown = (main_unknown || max_isr_unknown);

    mark_reachable("main");
    for (fn in functions) {
        if (fn ~ /^__vector_[0-9]+$/) {
            mark_reachable(fn);
        }
    }

    unknown_list = "";
    for (fn in functions) {
        if (reachable[fn] && !known[fn]) {
            if (unknown_list != "") {
                unknown_list = unknown_list ",";
            }
            unknown_list = unknown_list fn;
            combined_unknown = 1;
        }
    }

    indirect_list = "";
    for (fn in indirect_call) {
        if (indirect_call[fn] && reachable[fn]) {
            if (indirect_list != "") {
                indirect_list = indirect_list ",";
            }
            indirect_list = indirect_list fn;
            combined_unknown = 1;
        }
    }

    cycle_list = "";
    for (fn in cycle_node) {
        if (reachable[fn]) {
            if (cycle_list != "") {
                cycle_list = cycle_list ",";
            }
            cycle_list = cycle_list fn;
            combined_unknown = 1;
        }
    }

    confidence = "high";
    if (indirect_list != "" || cycle_list != "") {
        confidence = "low";
    } else if (unknown_list != "") {
        confidence = "medium";
    }

    print "STACK_MAIN_USED\t" main_total;
    print "STACK_MAIN_PATH\t" main_path;
    print "STACK_ISR_USED\t" max_isr_total;
    print "STACK_ISR_NAME\t" (max_isr_name == "" ? "(none)" : max_isr_name);
    print "STACK_ISR_PATH\t" (max_isr_path == "" ? "(none)" : max_isr_path);
    print "STACK_COMBINED_USED\t" combined_total;
    print "STACK_HAVE_ISR\t" (isr_count > 0 ? 1 : 0);
    print "STACK_UNKNOWN\t" (combined_unknown ? 1 : 0);
    print "STACK_CONFIDENCE\t" confidence;
    print "STACK_UNKNOWN_FUNCS\t" (unknown_list == "" ? "-" : unknown_list);
    print "STACK_INDIRECT_CALLERS\t" (indirect_list == "" ? "-" : indirect_list);
    print "STACK_CYCLE_NODES\t" (cycle_list == "" ? "-" : cycle_list);
}
' || true

    rm -f "$tmp_su"
}

text_used="$(section_size .text)"
data_used="$(section_size .data)"
bss_used="$(section_size .bss)"
noinit_used="$(section_size .noinit)"
bootloader_used="$(section_size .bootloader)"

flash_used=$((text_used + data_used + bootloader_used))
ram_static_used=$((data_used + bss_used + noinit_used))

device_macros="$(printf '#include <avr/io.h>\n' | avr-gcc -mmcu="$mcu" -dM -E -)"
flash_end_hex="$(awk '/^#define[[:space:]]+FLASHEND[[:space:]]+/ { print $3; exit }' <<<"$device_macros")"
ram_end_hex="$(awk '/^#define[[:space:]]+RAMEND[[:space:]]+/ { print $3; exit }' <<<"$device_macros")"
ram_start_hex="$(awk '/^#define[[:space:]]+RAMSTART[[:space:]]+/ { print $3; exit }' <<<"$device_macros")"

if [[ -z "$flash_end_hex" || -z "$ram_end_hex" || -z "$ram_start_hex" ]]; then
    echo "Unable to determine FLASHEND/RAMEND/RAMSTART for MCU: $mcu" >&2
    exit 1
fi

flash_total=$((flash_end_hex + 1))
ram_total=$((ram_end_hex - ram_start_hex + 1))

flash_free=$((flash_total - flash_used))
ram_stack_budget=$((ram_total - ram_static_used))

stack_main_used=0
stack_isr_used=0
stack_isr_name="(none)"
stack_main_path="(none)"
stack_isr_path="(none)"
stack_unknown=1
stack_confidence="low"
stack_unknown_funcs="-"
stack_indirect_callers="-"
stack_cycle_nodes="-"
have_su=0
for su in "${su_files[@]}"; do
    if [[ -f "$su" ]]; then
        have_su=1
        break
    fi
done

while IFS=$'\t' read -r key value; do
    case "$key" in
        STACK_MAIN_USED) stack_main_used="$value" ;;
        STACK_MAIN_PATH) stack_main_path="$value" ;;
        STACK_ISR_USED) stack_isr_used="$value" ;;
        STACK_ISR_NAME) stack_isr_name="$value" ;;
        STACK_ISR_PATH) stack_isr_path="$value" ;;
        STACK_COMBINED_USED) stack_used_est="$value" ;;
        STACK_UNKNOWN) stack_unknown="$value" ;;
        STACK_CONFIDENCE) stack_confidence="$value" ;;
        STACK_UNKNOWN_FUNCS) stack_unknown_funcs="$value" ;;
        STACK_INDIRECT_CALLERS) stack_indirect_callers="$value" ;;
        STACK_CYCLE_NODES) stack_cycle_nodes="$value" ;;
        STACK_ERROR) stack_unknown=1 ;;
    esac
done < <(estimate_stack "$elf_file" "${su_files[@]}")

: "${stack_used_est:=0}"

stack_free_est=$((ram_stack_budget - stack_used_est))
if ((stack_free_est < 0)); then
    stack_free_est=0
fi

ram_used_est=$((ram_static_used + stack_used_est))
ram_free_est=$((ram_total - ram_used_est))
if ((ram_free_est < 0)); then
    ram_free_est=0
fi

flash_pct=$((flash_used * 100 / flash_total))
ram_pct=$((ram_used_est * 100 / ram_total))

echo "Memory report for $elf_file ($mcu)"
echo "FLASH total: ${flash_total} bytes"
echo "FLASH used:  ${flash_used} bytes (${flash_pct}%)"
echo "FLASH free:  ${flash_free} bytes"
echo
echo "RAM total:   ${ram_total} bytes"
echo "RAM used:    ${ram_used_est} bytes (${ram_pct}%)"
echo "RAM free:    ${ram_free_est} bytes"
echo "  BSS used:      ${bss_used} bytes"
echo "  Data used:     ${data_used} bytes"
echo "  Noinit used:   ${noinit_used} bytes"
echo "  Stack budget:  ${ram_stack_budget} bytes"
echo "  Stack main*:   ${stack_main_used} bytes"
echo "  Stack ISR*:    ${stack_isr_used} bytes (${stack_isr_name})"
echo "  Stack used*:   ${stack_used_est} bytes"
echo "  Stack free*:   ${stack_free_est} bytes"
echo "  Stack conf*:   ${stack_confidence}"

if ((stack_unknown == 0)); then
    echo
    echo "* Stack used/free is a static max estimate from call-chain analysis (main + worst ISR)."
    echo "  Main path: ${stack_main_path}"
    echo "  ISR path:  ${stack_isr_path}"
elif ((have_su == 0)); then
    echo
    echo "* Stack estimate is low-confidence (.su file not found); reported stack uses fallback analysis."
else
    echo
    echo "* Stack estimate includes unresolved uncertainty."
    echo "  Main path: ${stack_main_path}"
    echo "  ISR path:  ${stack_isr_path}"
    if [[ "$stack_unknown_funcs" != "-" ]]; then
        echo "  Unknown stack functions: ${stack_unknown_funcs}"
    fi
    if [[ "$stack_indirect_callers" != "-" ]]; then
        echo "  Indirect-call sites: ${stack_indirect_callers}"
    fi
    if [[ "$stack_cycle_nodes" != "-" ]]; then
        echo "  Recursive/cyclic nodes: ${stack_cycle_nodes}"
    fi
fi