#!/usr/bin/env python3
"""
extract_cbor_uart.py - Extract UART serial output from CBOR hardware trace files.

Usage:
    python3 tools/extract_cbor_uart.py <file.cbor>              # print to stdout
    python3 tools/extract_cbor_uart.py <file.cbor> -o out.txt   # save to file
    python3 tools/extract_cbor_uart.py <file.cbor> -s           # summary only
"""

import cbor2
import sys
import argparse
import re
from pathlib import Path


def extract_uart_output(data, verbose=False):
    """Extract UART text from CBOR trace events.

    The CBOR trace format contains events where UART output appears as
    byte arrays within list records. We scan all byte payloads that
    contain printable ASCII and concatenate them.
    """
    output_parts = []
    total_bytes = 0
    uart_bytes = 0
    event_count = len(data) if isinstance(data, list) else 0

    for i, item in enumerate(data):
        if not isinstance(item, list):
            continue

        for elem in item:
            if not isinstance(elem, bytes):
                continue
            total_bytes += len(elem)

            # Filter: UART output is printable ASCII with CR/LF
            # (as opposed to binary memory dumps)
            try:
                text = elem.decode('ascii')
                # UART text has a high proportion of printable chars
                printable = sum(1 for c in text if c.isprintable() or c in '\r\n\t')
                ratio = printable / max(len(text), 1)
                if ratio > 0.5 and len(text) <= 1024:
                    uart_bytes += len(elem)
                    if verbose:
                        print(f"  Event[{i}] +{len(elem)} bytes: {text[:60].rstrip()}", 
                              file=sys.stderr)
                    output_parts.append(text)
            except (UnicodeDecodeError, ValueError):
                pass  # binary data, skip

    return ''.join(output_parts), uart_bytes, total_bytes, event_count


def parse_timing(text):
    """Extract timing statistics from UART output."""
    patterns = {
        'crc32':       r'MATMUL_CRC32=([0-9a-fA-F]+)',
        'prologue':    r'TIME_PROLOGUE=(\d+)',
        'loop':        r'TIME_LOOP=(\d+)',
        'loop_per':    r'TIME_LOOP_PER=(\d+)',
        'epilogue':    r'TIME_EPILOGUE=(\d+)',
        'total':       r'TIME_TOTAL=(\d+)',
        # Pipeline diagnostics (v2.1)
        'wait_ip':     r'PIPE_WAIT_IP=(\d+)',
        'wait_ip_avg': r'PIPE_WAIT_IP_AVG=(\d+)',
        'wait_dma0':   r'PIPE_WAIT_DMA0=(\d+)',
        'wait_dma0_avg': r'PIPE_WAIT_DMA0_AVG=(\d+)',
        'wait_dma1':   r'PIPE_WAIT_DMA1=(\d+)',
        'wait_dma1_avg': r'PIPE_WAIT_DMA1_AVG=(\d+)',
        'fire_dma':    r'PIPE_FIRE_DMA=(\d+)',
        'fire_ip':     r'PIPE_FIRE_IP=(\d+)',
        'stall_ip':    r'PIPE_STALL_IP=(\d+)',
        'stall_dma0':  r'PIPE_STALL_DMA0=(\d+)',
        'stall_dma1':  r'PIPE_STALL_DMA1=(\d+)',
        'sum_waits':   r'PIPE_SUM_WAITS=(\d+)',
    }
    results = {}
    for key, pattern in patterns.items():
        m = re.search(pattern, text)
        if m:
            val = m.group(1)
            try:
                results[key] = int(val, 16) if 'crc' in key else int(val)
            except ValueError:
                results[key] = val
    return results


def main():
    parser = argparse.ArgumentParser(
        description='Extract UART serial output from CBOR hardware trace files.')
    parser.add_argument('file', type=Path, help='CBOR trace file (.cbor)')
    parser.add_argument('-o', '--output', type=Path,
                        help='Write output to file (default: stdout)')
    parser.add_argument('-s', '--summary', action='store_true',
                        help='Print timing summary only, not full UART text')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Verbose event-by-event output to stderr')
    args = parser.parse_args()

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

    # Load CBOR
    with open(args.file, 'rb') as f:
        data = cbor2.load(f)

    # Extract UART
    uart_text, uart_bytes, total_bytes, event_count = extract_uart_output(
        data, verbose=args.verbose)

    # Summary to stderr
    print(f"File: {args.file.name}", file=sys.stderr)
    print(f"  Events: {event_count}, Total bytes: {total_bytes}, "
          f"UART bytes: {uart_bytes}", file=sys.stderr)

    # Parse timing
    timing = parse_timing(uart_text)
    if timing:
        print(f"  MATMUL_CRC32: {timing.get('crc32', 'N/A')}", file=sys.stderr)
        print(f"  TIME_TOTAL:   {timing.get('total', 'N/A')} ns "
              f"({timing.get('total', 0)/1e6:.2f} ms)", file=sys.stderr)
        if 'loop_per' in timing:
            print(f"  LOOP_PER:     {timing.get('loop_per', 'N/A')} ns/iter",
                  file=sys.stderr)

        # Pipeline profile
        pipe_keys = ['wait_ip', 'wait_dma0', 'wait_dma1', 'fire_dma', 'fire_ip',
                     'stall_ip', 'stall_dma0', 'stall_dma1']
        pipe_data = {k: timing.get(k) for k in pipe_keys if k in timing}
        if pipe_data:
            print(f"  --- Pipeline ---", file=sys.stderr)
            for k in pipe_keys:
                if k in timing:
                    label = k.replace('_', ' ').upper()
                    val = timing[k]
                    print(f"  {label:20s} = {val}", file=sys.stderr)

    # Output
    if args.summary:
        print(f"\nSummary saved. Use without -s to see full UART output.",
              file=sys.stderr)
    elif args.output:
        args.output.write_text(uart_text)
        print(f"UART output saved to: {args.output} ({len(uart_text)} bytes)",
              file=sys.stderr)
    else:
        sys.stdout.write(uart_text)


if __name__ == '__main__':
    main()