#!/usr/bin/env python3
"""firmware_overview.py -- Quick structural overview of ADSP-219x firmware.

Analyzes a PM binary and produces a human-readable report:
  - Memory map (code vs data vs empty regions)
  - Function-like blocks (code between RTS/JUMP boundaries)
  - DSP kernels (DO UNTIL loops with MAC operations)
  - I/O configuration (peripheral register writes)
  - All immediate constants loaded into registers
  - Cross-reference: which addresses are referenced

Usage:
    python3 firmware_overview.py firmware_pm.bin
    python3 firmware_overview.py firmware_pm.bin --dm firmware_dm.bin
"""

import argparse
import os
import struct
import sys


# ADSP-2191 I/O register names (partial, most common)
IO_NAMES = {
    0x00: "SYSCTRL", 0x01: "SYSTAT", 0x02: "IOPG",
    0x04: "BWCTRL", 0x05: "BWSTAT",
    0x10: "PICR0", 0x11: "PICR1", 0x12: "PICR2", 0x13: "PICR3",
    0x14: "IMASK0", 0x15: "IMASK1",
    0x20: "TCTL", 0x21: "TCOUNT", 0x22: "TPERIOD", 0x23: "TSCALE",
    0x40: "SPORT0_CTL", 0x41: "SPORT0_DIV",
    0x42: "SPORT0_MCTL", 0x43: "SPORT0_CS0",
    0x44: "SPORT0_TX", 0x45: "SPORT0_RX",
    0x46: "SPORT0_STAT",
    0x60: "SPORT1_CTL", 0x61: "SPORT1_DIV",
    0x62: "SPORT1_MCTL", 0x63: "SPORT1_CS0",
    0x64: "SPORT1_TX", 0x65: "SPORT1_RX",
    0x66: "SPORT1_STAT",
    0x80: "SPICTL", 0x81: "SPISTAT", 0x82: "SPIFLG",
    0x83: "SPIBAUD", 0x84: "SPIRX", 0x85: "SPITX",
    0xA0: "UART_TX", 0xA1: "UART_RX", 0xA2: "UART_CTL",
    0xA3: "UART_STAT", 0xA4: "UART_DIV",
    0xC0: "FLAGD", 0xC1: "FLAGP", 0xC2: "FLAGS",
    0xC3: "FLAGC",
}

REG0 = ["AX0", "AX1", "MX0", "MX1", "AY0", "AY1", "MY0", "MY1",
        "MR2", "SR2", "AR", "SI", "MR1", "SR1", "MR0", "SR0"]
REG1 = ["I0", "I1", "I2", "I3", "M0", "M1", "M2", "M3",
        "L0", "L1", "L2", "L3", "IMASK", "IRPTL", "ICNTL", "STACKA"]
REG2 = ["I4", "I5", "I6", "I7", "M4", "M5", "M6", "M7",
        "L4", "L5", "L6", "L7", "RES", "RES", "CNTR", "LPSTACKA"]


def read_pm(path):
    """Read 24-bit PM words from packed binary."""
    data = open(path, "rb").read()
    words = []
    for i in range(0, len(data) - 2, 3):
        w = (data[i] << 16) | (data[i + 1] << 8) | data[i + 2]
        words.append(w)
    return words


def classify_word(w):
    """Return a rough classification of a PM word."""
    if w == 0:
        return "null"
    b23_22 = (w >> 22) & 3
    if b23_22 == 3:
        return "compute_multi"  # Type 1
    b23_16 = (w >> 16) & 0xFF
    if b23_16 == 0x0A:
        return "rts"
    if (w >> 18) == 0x07:
        return "jump"
    if (w >> 19) == 0x03:
        return "jump_cond"
    if b23_16 == 0x16:
        return "do_until"
    if b23_16 == 0x0B:
        return "jump_indirect"
    if b23_16 == 0x05:
        return "ljump"
    return "code"


def find_constants(words):
    """Extract all immediate constant loads."""
    constants = []
    for i, w in enumerate(words):
        addr = i * 3
        b23_22 = (w >> 22) & 3
        b21_20 = (w >> 20) & 3

        # Type 6: Dreg = Imm16
        if b23_22 == 1 and b21_20 == 0:
            val = (w >> 4) & 0xFFFF
            reg = REG0[w & 0xF]
            constants.append((addr, reg, val, "dreg"))

        # Type 7 REG1: Reg1 = Imm16
        elif b23_22 == 1 and b21_20 == 1:
            val = (w >> 4) & 0xFFFF
            reg = REG1[w & 0xF]
            constants.append((addr, reg, val, "reg1"))

        # Type 7 REG2: Reg2 = Imm16
        elif b23_22 == 0 and b21_20 == 3:
            val = (w >> 4) & 0xFFFF
            reg = REG2[w & 0xF]
            constants.append((addr, reg, val, "reg2"))

        # Type 33: Reg3 = Data12
        elif (w >> 16) == 0x10:
            val = (w >> 4) & 0xFFF
            constants.append((addr, "REG3", val, "short"))

        # Type 34: IO write
        if (w >> 16) == 0x06 and ((w >> 15) & 1):
            d = (w >> 12) & 1
            io_addr = (((w >> 13) & 3) << 8) | ((w >> 4) & 0xFF)
            dreg = REG0[w & 0xF]
            io_name = IO_NAMES.get(io_addr, f"IO_0x{io_addr:03X}")
            if d:
                constants.append((addr, io_name, None,
                                  f"io_write({dreg})"))
            else:
                constants.append((addr, dreg, None,
                                  f"io_read({io_name})"))

    return constants


def find_do_loops(words):
    """Find DO UNTIL loops and their bodies."""
    loops = []
    for i, w in enumerate(words):
        if (w >> 16) == 0x16:
            rel = (w >> 4) & 0xFFF
            term = w & 0xF
            cond = ["EQ", "NE", "GT", "LE", "LT", "GE",
                    "AV", "NOT AV", "AC", "NOT AC",
                    "SWCOND", "NOT SWCOND", "MV", "NOT MV",
                    "NOT CE", "TRUE"][term]
            srel = rel if rel < 0x800 else rel - 0x1000
            target = i * 3 + srel * 3
            # Count MAC ops in loop body
            mac_count = 0
            for j in range(i + 1, min(i + srel + 1, len(words))):
                if (words[j] >> 22) & 3 == 3:  # Type 1
                    amf = (words[j] >> 13) & 0x1F
                    if amf < 16:  # MAC operation
                        mac_count += 1
            loops.append((i * 3, target, cond, srel, mac_count))
    return loops


def memory_map(words, chunk_size=64):
    """Build a coarse memory map."""
    regions = []
    i = 0
    while i < len(words):
        end = min(i + chunk_size, len(words))
        chunk = words[i:end]
        n_null = sum(1 for w in chunk if w == 0)
        n_multi = sum(1 for w in chunk
                      if (w >> 22) & 3 == 3)
        n_jump = sum(1 for w in chunk
                     if classify_word(w) in
                     ("jump", "jump_cond", "rts", "do_until"))

        if n_null > len(chunk) * 0.9:
            kind = "empty"
        elif n_multi > len(chunk) * 0.3:
            kind = "dsp_kernel"
        elif n_jump > len(chunk) * 0.1:
            kind = "control_flow"
        else:
            kind = "code/data"

        if regions and regions[-1][2] == kind:
            regions[-1] = (regions[-1][0], end * 3, kind)
        else:
            regions.append((i * 3, end * 3, kind))
        i = end
    return regions


def analyze_pm(path, dm_path=None):
    """Full PM analysis."""
    words = read_pm(path)
    size = len(words) * 3

    print(f"=== PM Firmware Overview ===")
    print(f"File: {path}")
    print(f"Size: {size} bytes ({len(words)} instructions)")
    print()

    # Memory map
    regions = memory_map(words)
    print(f"--- Memory Map ---")
    for start, end, kind in regions:
        label = {"empty": "EMPTY (null)",
                 "dsp_kernel": "DSP KERNEL (MAC-heavy)",
                 "control_flow": "CONTROL FLOW (jumps)",
                 "code/data": "CODE / DATA"}[kind]
        print(f"  0x{start:06X}-0x{end:06X}"
              f"  ({(end-start)//3:4d} words)  {label}")
    print()

    # Entry point
    if words:
        w0 = words[0]
        if (w0 >> 18) == 0x07:
            rel = ((w0 >> 4) & 0x3FFF) | ((w0 & 3) << 14)
            srel = rel if rel < 0x8000 else rel - 0x10000
            target = srel * 3
            print(f"--- Entry Point ---")
            print(f"  Reset vector: JUMP 0x{target:06X}")
            print()
        elif w0 == 0:
            print(f"--- Entry Point ---")
            print(f"  Reset vector: NOP (code starts at 0x000003)")
            print()

    # DO UNTIL loops (DSP kernels)
    loops = find_do_loops(words)
    if loops:
        print(f"--- DSP Loops ({len(loops)} found) ---")
        for addr, end, cond, length, macs in loops:
            kind = ""
            if macs > 0:
                kind = f" [MAC kernel, {macs} multiply-accumulate ops]"
            elif cond == "NOT CE":
                kind = " [counter loop]"
            print(f"  0x{addr:06X}: DO 0x{end:06X} UNTIL {cond}"
                  f" ({length} words){kind}")
        print()

    # Constants
    constants = find_constants(words)
    if constants:
        # Group by type
        reg_loads = [(a, r, v, t) for a, r, v, t in constants
                     if t in ("dreg", "reg1", "reg2", "short")]
        io_ops = [(a, r, v, t) for a, r, v, t in constants
                  if t.startswith("io_")]

        if reg_loads:
            print(f"--- Constants ({len(reg_loads)} register loads) ---")
            # Show unique values
            seen = {}
            for addr, reg, val, _ in reg_loads:
                key = (reg, val)
                if key not in seen:
                    seen[key] = []
                seen[key].append(addr)

            for (reg, val), addrs in sorted(seen.items(),
                                            key=lambda x: x[0][1]):
                sval = val - 0x10000 if val >= 0x8000 else val
                locs = ", ".join(f"0x{a:06X}" for a in addrs[:3])
                more = f" +{len(addrs)-3} more" if len(addrs) > 3 else ""
                print(f"  {reg:8s} = 0x{val:04X}"
                      f" ({sval:6d})  @ {locs}{more}")
            print()

        if io_ops:
            print(f"--- I/O Operations ({len(io_ops)} found) ---")
            for addr, reg, _, typ in io_ops:
                print(f"  0x{addr:06X}: {typ}")
            print()

    # DM analysis if provided
    if dm_path and os.path.isfile(dm_path):
        print(f"--- DM Cross-Reference ---")
        dm_data = open(dm_path, "rb").read()
        dm_size = len(dm_data)
        print(f"  DM file: {dm_path} ({dm_size} bytes,"
              f" {dm_size//2} words)")

        # Find all I-register loads and show what's at those
        # DM addresses
        for addr, reg, val, typ in constants:
            if reg.startswith("I") and reg[1:].isdigit() \
                    and typ in ("reg1", "reg2"):
                dm_byte = val * 2
                if dm_byte + 20 <= dm_size:
                    sample = []
                    for j in range(10):
                        w = struct.unpack_from(
                            ">H", dm_data, dm_byte + j * 2)[0]
                        sample.append(f"0x{w:04X}")
                    print(f"  {reg} = 0x{val:04X} -> DM content:"
                          f" [{', '.join(sample[:5])} ...]")
        print()

    # Summary
    n_code = sum(1 for w in words if w != 0)
    n_mac = sum(1 for w in words if (w >> 22) & 3 == 3)
    n_jump = sum(1 for w in words
                 if classify_word(w) in
                 ("jump", "jump_cond", "jump_indirect"))
    n_rts = sum(1 for w in words
                if classify_word(w) == "rts")

    print(f"--- Summary ---")
    print(f"  Total words:    {len(words)}")
    print(f"  Non-zero:       {n_code}"
          f" ({n_code*100//max(len(words),1)}%)")
    print(f"  Multifunction:  {n_mac} (MAC/ALU + memory)")
    print(f"  Jumps/Calls:    {n_jump}")
    print(f"  Returns:        {n_rts}")
    print(f"  DO loops:       {len(loops)}")
    print(f"  I/O accesses:   {len(io_ops) if constants else 0}")


def main():
    parser = argparse.ArgumentParser(
        description="Quick overview of ADSP-219x PM firmware.")
    parser.add_argument("file", help="PM binary file (packed 3-byte)")
    parser.add_argument("--dm", help="Optional DM binary for"
                        " cross-reference")
    args = parser.parse_args()

    if not os.path.isfile(args.file):
        print(f"Error: {args.file} not found", file=sys.stderr)
        sys.exit(1)

    analyze_pm(args.file, args.dm)


if __name__ == "__main__":
    main()