#!/usr/bin/env python3
"""
UART Interactive Terminal — CCSDS122 register & DDR access

Commands:
  w  <addr> <data>          Write 32-bit register (e.g. w 0x80000000 0x01)
  r  <addr>                 Read  32-bit register (e.g. r 0x80000010)
  aw <addr> <len> <file>    Write DDR from binary file
  aw <addr> hex:<hexstr>    Write DDR from hex string (e.g. aw 0 hex:deadbeef)
  ar <addr> <len>           Read  DDR, print hex dump
  ar <addr> <len> <file>    Read  DDR, save to file
  help                      Show this help
  quit                      Exit

Protocol:
  Reg Write: 'W' [4B addr BE] [4B data BE]
  Reg Read:  'R' [4B addr BE]  →  4B response
  DDR Write: 'A' [4B addr BE] [4B len BE] [data bytes]
  DDR Read:  'D' [4B addr BE] [4B len BE]  →  len bytes response
"""

import serial
import struct
import sys
import os

# ---- UART Command Primitives ------------------------------------------------

def send_byte_raw(ser, b):
    """Send a single raw byte (for escape/test)."""
    ser.write(bytes([b]))

def reg_write(ser, addr, data):
    """AXI-Lite register write."""
    cmd = b'W' + struct.pack('>I', addr) + struct.pack('>I', data)
    ser.write(cmd)
    print(f"  W 0x{addr:08X} = 0x{data:08X}  ({data})")

def reg_read(ser, addr):
    """AXI-Lite register read, returns 32-bit value."""
    cmd = b'R' + struct.pack('>I', addr)
    ser.write(cmd)
    raw = ser.read(4)
    if len(raw) < 4:
        raise IOError(f"Reg read @ 0x{addr:08X}: timeout, got {len(raw)} bytes")
    return struct.unpack('>I', raw)[0]

def ddr_write(ser, addr, data_bytes):
    """Write data to DDR at specific address via 'A' command."""
    cmd = b'A' + struct.pack('>I', addr) + struct.pack('>I', len(data_bytes))
    ser.write(cmd + data_bytes)

def ddr_read(ser, addr, length, timeout=10):
    """Read data from DDR via 'D' command, returns bytes.
       内部自动补齐到 16B 对齐, 避免末拍不满导致字节排序错乱."""
    aligned = ((length + 15) // 16) * 16
    cmd = b'D' + struct.pack('>I', addr) + struct.pack('>I', aligned)
    ser.write(cmd)
    ser.timeout = min(timeout, max(10, aligned * 10 / 2000000 * 3))
    raw = ser.read(aligned)
    if len(raw) < length:
        print(f"  WARNING: expected {length} bytes, got {len(raw)}")
    return raw[:length]

# ---- Hex Dump ---------------------------------------------------------------

def hexdump(data, addr=0, show_all=False):
    """Print hex dump of data bytes."""
    if not show_all and len(data) > 256:
        for i in range(0, 256, 16):
            chunk = data[i:i+16]
            hexstr = ' '.join(f'{b:02X}' for b in chunk)
            ascii_str = ''.join(chr(b) if 32 <= b < 127 else '.' for b in chunk)
            print(f"  {addr+i:08X}  {hexstr:<48s}  |{ascii_str}|")
        print(f"  ... ({len(data) - 256} more bytes)")
    else:
        for i in range(0, len(data), 16):
            chunk = data[i:i+16]
            hexstr = ' '.join(f'{b:02X}' for b in chunk)
            ascii_str = ''.join(chr(b) if 32 <= b < 127 else '.' for b in chunk)
            print(f"  {addr+i:08X}  {hexstr:<48s}  |{ascii_str}|")

# ---- Register Name Lookup ---------------------------------------------------

REG_NAMES = {
    0x80000000: "CTRL",
    0x80000004: "IMG_ADDR",
    0x80000008: "IMG_LEN",
    0x8000000C: "CMPR_ADDR",
    0x80000010: "STATUS",
    0x80000014: "CMPR_LEN",
    0x80000018: "VERSION",
    0x8000001C: "STORAGE_MODE",
    0x80000020: "MEM_BLOCK_SZ",
    0x80000024: "PAN_BYTES",
    0x80000028: "B1_BYTES",
    0x8000002C: "B2_BYTES",
    0x80000030: "B3_BYTES",
    0x80000034: "B4_BYTES",
    0x80000038: "TOTAL_BYTES",
    0x8000003C: "IMG_PER_REGION",
}

def reg_name(addr):
    for base, name in REG_NAMES.items():
        if addr == base:
            return name
    return ""

# ---- Command Parser ---------------------------------------------------------

def parse_int(s):
    """Parse hex (0x...) or decimal integer."""
    s = s.strip()
    if s.lower().startswith('0x'):
        return int(s, 16)
    return int(s)

def cmd_w(ser, parts):
    """w <addr> <data>"""
    if len(parts) < 3:
        print("  Usage: w <addr> <data>")
        return
    addr = parse_int(parts[1])
    data = parse_int(parts[2])
    name = reg_name(addr)
    extra = f" ({name})" if name else ""
    print(f"  Write 0x{addr:08X}{extra}")
    reg_write(ser, addr, data)

def cmd_r(ser, parts):
    """r <addr>"""
    if len(parts) < 2:
        print("  Usage: r <addr>")
        return
    addr = parse_int(parts[1])
    name = reg_name(addr)
    extra = f" ({name})" if name else ""
    print(f"  Read 0x{addr:08X}{extra}")
    val = reg_read(ser, addr)
    print(f"  →  0x{val:08X}  ({val})")
    # Decode STATUS register
    if addr == 0x80000010:
        print(f"       busy={val & 1},  done={(val>>1) & 1}")
    elif addr == 0x80000000:
        print(f"       enable={val & 1},  bpp={(val>>1) & 7}")

def cmd_aw(ser, parts):
    """aw <addr> <len> <file>                    Write file to DDR at addr
       aw <addr> <len> hex:<hexstr>             Write hex string to DDR at addr
       aw <addr> <len> --loop <rows> <stride> <file>   Strided write from file"""
    if len(parts) < 3:
        print("  Usage: aw <addr> <len> <file>")
        print("         aw <addr> <len> hex:<hexstr>")
        print("         aw <addr> <len> --loop <rows> <stride> <file>")
        return
    addr = parse_int(parts[1])
    length = parse_int(parts[2])

    # ---- Loop mode: --loop <rows> <stride> <file> --------------------------
    loop_mode = "--loop" in parts
    if loop_mode:
        loop_idx = parts.index("--loop")
        loop_rows = int(parts[loop_idx+1])
        loop_stride = int(parts[loop_idx+2])
        fname_idx = loop_idx + 3
        if len(parts) <= fname_idx:
            print("  ERROR: missing file path for --loop mode")
            return
        fpath = os.path.expandvars(os.path.expanduser(parts[fname_idx]))
        if not os.path.exists(fpath):
            print(f"  ERROR: file not found: {fpath}")
            return
        with open(fpath, 'rb') as f:
            file_data = f.read()

        total_written = 0
        print(f"  Write DDR 0x{addr:08X}, {length}B/row, {loop_rows} rows, stride={loop_stride}")
        for row in range(loop_rows):
            raddr = addr + row * loop_stride
            src_start = row * length
            src_end = src_start + length
            if src_start >= len(file_data):
                print(f"  WARNING: file exhausted at row {row}")
                break
            chunk = file_data[src_start:src_end]
            if len(chunk) < length:
                chunk = chunk + bytes(length - len(chunk))
            ddr_write(ser, raddr, chunk)
            total_written += len(chunk)
            if row % 64 == 0 or row == loop_rows - 1:
                print(f"  row {row:4d}/{loop_rows}  @0x{raddr:08X}  {len(chunk)}B")
        print(f"  Loop done: {loop_rows} rows, {total_written} bytes written")
        return

    # ---- Normal mode ---------------------------------------------------------
    src = parts[2]

    if src.lower().startswith('hex:'):
        hexstr = src[4:].replace(' ', '')
        data = bytes.fromhex(hexstr)
    else:
        fpath = os.path.expandvars(os.path.expanduser(src))
        if not os.path.exists(fpath):
            print(f"  ERROR: file not found: {fpath}")
            return
        with open(fpath, 'rb') as f:
            data = f.read()

    print(f"  Write DDR 0x{addr:08X}, {len(data)} bytes")
    ddr_write(ser, addr, data)
    print(f"  Done")

def cmd_ar(ser, parts):
    """ar <addr> <len> [file]   or   ar <addr> <len> --loop <rows> <stride> [file_prefix]"""
    if len(parts) < 3:
        print("  Usage: ar <addr> <len> [file]")
        print("         ar <addr> <len> --loop <rows> <stride> [file_prefix]")
        return
    addr = parse_int(parts[1])
    length = parse_int(parts[2])

    # ---- Loop mode: --loop <rows> <stride> ----------------------------------
    loop_mode = "--loop" in parts
    if loop_mode:
        loop_idx = parts.index("--loop")
        loop_rows = int(parts[loop_idx+1])
        loop_stride = int(parts[loop_idx+2])
        fname_idx = loop_idx + 3
        fpath = parts[fname_idx] if len(parts) > fname_idx else None
        all_data = bytearray()
        total_read = 0
        for row in range(loop_rows):
            raddr = addr + row * loop_stride
            data = ddr_read(ser, raddr, length, timeout=max(10, length * 10 / 2000000 * 4))
            all_data.extend(data)
            total_read += len(data)
            if row % 64 == 0 or row == loop_rows - 1:
                print(f"  row {row:4d}/{loop_rows}  @0x{raddr:08X}  {len(data)}B")
        if fpath:
            fpath = os.path.expandvars(os.path.expanduser(fpath))
            with open(fpath, 'wb') as f:
                f.write(all_data)
            print(f"  Saved {total_read} bytes to {fpath}")
        else:
            hexdump(all_data, addr, show_all=(len(all_data) <= 1024))
        print(f"  Loop done: {loop_rows} rows, {total_read} bytes total")
        return

    # ---- Normal mode ---------------------------------------------------------
    fpath = parts[3] if len(parts) > 3 else None
    print(f"  Read DDR 0x{addr:08X}, {length} bytes")
    data = ddr_read(ser, addr, length, timeout=max(10, length * 10 / 2000000 * 4))

    if fpath:
        fpath = os.path.expandvars(os.path.expanduser(fpath))
        with open(fpath, 'wb') as f:
            f.write(data)
        print(f"  Saved {len(data)} bytes to {fpath}")
    else:
        hexdump(data, addr, show_all=(length <= 1024))

def cmd_help():
    print("""Commands:
  w  <addr> <data>          Write 32-bit register
  r  <addr>                 Read  32-bit register
  aw <addr> <len> <file>          Write DDR from binary file
  aw <addr> <len> hex:<hexstr>    Write DDR from hex string
  ar <addr> <len>           Read  DDR, print hex dump
  aw <addr> <len> --loop <rows> <stride> <file>    Strided write
  ar <addr> <len> --loop <rows> <stride> [file]    Strided read
  ar <addr> <len> <file>    Read  DDR, save to file
  dump                      Show all registers
  help                      Show this help
  quit / exit               Exit

Addresses:
  0x00000000 - DDR (raw image data)
  0x40000000 - DDR (compressed output)
  0x80000000+  - axil_slave registers

Examples:
  r 0x80000018               Read VERSION
  aw 0x03000000 512 --loop 256 3072 b1.bin   Strided write B1 region
  w 0x80000000 0x01          Write CTRL = enable
  ar 0x00000000 512 --loop 256 12288 pan.bin  Strided read PAN region
  r 0x80000010               Read STATUS
  aw 0x00000000 test.bin     Write test.bin to DDR addr 0
  ar 0x40000000 1024         Read 1024 bytes from DDR
  ar 0x40000000 4096 out.bin Read DDR, save to file
""")

def cmd_dump(ser):
    """Read all known registers."""
    print(f"\n  {'Addr':<12} {'Name':<16} {'Value':<12} {'Dec'}")
    print(f"  {'-'*12} {'-'*16} {'-'*12} {'-'*10}")
    for addr, name in sorted(REG_NAMES.items()):
        val = reg_read(ser, addr)
        print(f"  0x{addr:08X}  {name:<16s}  0x{val:08X}  ({val})")
        if addr == 0x80000010:
            print(f"    {'':>12s}  {'':>16s}    busy={val & 1}, done={(val>>1) & 1}")
        elif addr == 0x80000000:
            print(f"    {'':>12s}  {'':>16s}    enable={val & 1}, bpp={(val>>1) & 7}")

# =============================================================================

def main():
    import argparse
    parser = argparse.ArgumentParser(description="CCSDS122 UART interactive terminal")
    parser.add_argument("--port", default="COM3", help="Serial port")
    parser.add_argument("--baud", type=int, default=2000000, help="Baud rate")
    args = parser.parse_args()

    print(f"Opening {args.port} @ {args.baud}...")
    ser = serial.Serial(args.port, args.baud, timeout=10, write_timeout=10)

    # Let FPGA UART RX stabilize
    import time as _time
    ser.rts = 0; _time.sleep(0.002); ser.rts = 1; _time.sleep(0.005)

    print("Connected. Type 'help' for commands.\n")

    try:
        while True:
            try:
                line = input("uart> ").strip()
            except (EOFError, KeyboardInterrupt):
                print("\nquit")
                break

            if not line:
                continue

            parts = line.split()
            cmd = parts[0].lower()

            if cmd in ('quit', 'exit', 'q'):
                break
            elif cmd == 'help':
                cmd_help()
            elif cmd == 'w':
                cmd_w(ser, parts)
            elif cmd == 'r':
                cmd_r(ser, parts)
            elif cmd == 'aw':
                cmd_aw(ser, parts)
            elif cmd == 'ar':
                cmd_ar(ser, parts)
            elif cmd == 'dump':
                cmd_dump(ser)
            else:
                print(f"  Unknown command: {cmd}. Type 'help'.")
    finally:
        ser.close()
        print("Closed.")


if __name__ == "__main__":
    main()