#!/usr/bin/env python3

## Imports

import sys
import re
import argparse

## Constants

VERSION = '1.0.0'

# https://www.raspberrypi.com/documentation/computers/raspberry-pi.html#raspberry-pi-revision-codes
CPUINFO_PATH = '/proc/cpuinfo'
CPUINFO_RE   = r'^Revision\s*:\s*(.*)$'

FIELD_RE = r'((.)\2*)'
FIELDS = [
    # Old-style.
    # https://www.raspberrypi.com/documentation/computers/raspberry-pi.html#old-style-revision-codes
    # https://forums.raspberrypi.com/viewtopic.php?p=176865#p177014
    'uuuuuuuWFuuuuuuuIIIIIIIIIIIIIIII',
    # New-style.
    # https://www.raspberrypi.com/documentation/computers/raspberry-pi.html#new-style-revision-codes
    'NOQuuuWuFMMMCCCCPPPPTTTTTTTTRRRR',
]

SYMBOL_UNUSED   = 'u'
SYMBOL_NEW_FLAG = 'F'
SYMBOLS = {
    'I':       [["Type", "Revision", "Memory size", "Manufacturer"], {
        0x0002: ["B",    "1.0",      "256MB",       "Egoman"      ],
        0x0003: ["B",    "1.0",      "256MB",       "Egoman"      ],
        0x0004: ["B",    "2.0",      "256MB",       "Sony UK"     ],
        0x0005: ["B",    "2.0",      "256MB",       "Qisda"       ],
        0x0006: ["B",    "2.0",      "256MB",       "Egoman"      ],
        0x0007: ["A",    "2.0",      "256MB",       "Egoman"      ],
        0x0008: ["A",    "2.0",      "256MB",       "Sony UK"     ],
        0x0009: ["A",    "2.0",      "256MB",       "Qisda"       ],
        0x000d: ["B",    "2.0",      "512MB",       "Egoman"      ],
        0x000e: ["B",    "2.0",      "512MB",       "Sony UK"     ],
        0x000f: ["B",    "2.0",      "512MB",       "Egoman"      ],
        0x0010: ["B+",   "1.2",      "512MB",       "Sony UK"     ],
        0x0011: ["CM1",  "1.0",      "512MB",       "Sony UK"     ],
        0x0012: ["A+",   "1.1",      "256MB",       "Sony UK"     ],
        0x0013: ["B+",   "1.2",      "512MB",       "Embest"      ],
        0x0014: ["CM1",  "1.0",      "512MB",       "Embest"      ],
        0x0015: ["A+",   "1.1",      "256MB/512MB", "Embest"      ],
    }],
    'R':       [["Revision"], {
        rev:    [f"1.{rev}"] for rev in range(2**len('RRRR'))
    }],
    'T':       [["Type"], {
        0x00:   ["A"],
        0x01:   ["B"],
        0x02:   ["A+"],
        0x03:   ["B+"],
        0x04:   ["2B"],
        0x05:   ["Alpha (early prototype)"],
        0x06:   ["CM1"],
        0x08:   ["3B"],
        0x09:   ["Zero"],
        0x0a:   ["CM3"],
        0x0c:   ["Zero W"],
        0x0d:   ["3B+"],
        0x0e:   ["3A+"],
        0x0f:   ["Internal use only"],
        0x10:   ["CM3+"],
        0x11:   ["4B"],
        0x12:   ["Zero 2 W"],
        0x13:   ["400"],
        0x14:   ["CM4"],
        0x15:   ["CM4S"],
        0x16:   ["Internal use only"],
        0x17:   ["5"],
        0x18:   ["CM5"],
        0x19:   ["500"],
        0x1a:   ["CM5 Lite"],
    }],
    'P':       [["Processor"], {
        0x0:    ["BCM2835"],
        0x1:    ["BCM2836"],
        0x2:    ["BCM2837"],
        0x3:    ["BCM2711"],
        0x4:    ["BCM2712"],
    }],
    'C':       [["Manufacturer"], {
        0x0:    ["Sony UK"],
        0x1:    ["Egoman"],
        0x2:    ["Embest"],
        0x3:    ["Sony Japan"],
        0x4:    ["Embest"],
        0x5:    ["Stadium"],
    }],
    'M':       [["Memory size"], {
        0x0:    ["256MB"],
        0x1:    ["512MB"],
        0x2:    ["1GB"],
        0x3:    ["2GB"],
        0x4:    ["4GB"],
        0x5:    ["8GB"],
        0x6:    ["16GB"],
    }],
    'F':       [["New flag"], {
        0x0:    ["Old-style revision"],
        0x1:    ["New-style revision"],
    }],
    'W':       [["Warranty bit"], {
        0x0:    ["Warranty is intact"],
        0x1:    ["Warranty has been voided by overclocking"],
    }],
    'Q':       [["OTP Read"], {
        0x0:    ["OTP reading allowed"],
        0x1:    ["OTP reading disallowed"],
    }],
    'O':      [["OTP Program"], {
        0x0:   ["OTP programming allowed"],
        0x1:   ["OTP programming disallowed"],
    }],
    'N':      [["Overvoltage"], {
        0x0:   ["Overvoltage allowed"],
        0x1:   ["Overvoltage disallowed"],
    }],
}

## Helpers

def error(msg):
    print(msg, file=sys.stderr)
    exit(1)

def unexpected(name, value):
    return name, f"<UNEXPECTED BITS {value:b}>"

def parse(code):
    new_flag = (code >> FIELDS[0][::-1].index(SYMBOL_NEW_FLAG)) & 1
    fields   = FIELDS[new_flag]
    for field, symbol in re.findall(FIELD_RE, fields[::-1]):
        length = len(field)
        value  = code & ((1<<length)-1)
        code   = code >> length
        if symbol == SYMBOL_UNUSED:
            if value != 0:
                yield unexpected("Unused", value)
            continue
        names, values = SYMBOLS[symbol]
        descs = values.get(value)
        if not descs:
            yield unexpected(names[0], value)
            continue
        yield from zip(names, descs)

## Main

def main():
    # Parse arguments.
    parser = argparse.ArgumentParser(description="""
        Parse Raspberry Pi revision codes.
    """);
    parser.add_argument('--version', action='version', version=VERSION)
    parser.add_argument('<code>', nargs='?', help="""
        A hexadecimal revision code to parse. If not given, '/proc/cpuinfo'
        will be used to get the revision code of the current device.
    """)
    args = vars(parser.parse_args())
    # Get code string from argument.
    code = args['<code>']
    # Get code string from system.
    if not code:
        try:
            with open(CPUINFO_PATH) as file:
                match = re.search(CPUINFO_RE, file.read(), re.MULTILINE)
        except IOError:
            error(f"Could not open '{CPUINFO_PATH}'.")
        if not match:
            error(f"Could not find revision code in '{CPUINFO_PATH}'.")
        code = match.group(1)
    # Parse code string.
    try:
        code = int(code, 16)
    except ValueError:
        error(f"Could not parse revision code {code} as a hexadecimal number.")
    # Parse code.
    names_descs = list(parse(code))
    # Print.
    width = max(len(name) for name, _ in names_descs)
    for name, desc in names_descs:
        print(f"{name:{width}} : {desc}")

if __name__ == '__main__':
    main()
