from __future__ import annotations
import argparse, os, os.path, struct, sys
from typing import *

def name2int(s: Union[str, bytes]):
    if isinstance(s, str): return ord(s[0]) + ord(s[1])*256
    else:
        assert isinstance(s, bytes)
        return s[0] + s[1]*256

# supported section types:
# * SN: SFRs
# * VN: vectors
# * MI: memory layout
# * ES: 2nd SFR
# * MD: probably "metadata", version & timestam
#
# can be parsed:
# * SC: something with SFRs
# * FL: flash stuff
# * OD: debug stuff
#
# I know about their existence:
# * CG, XD, CT, QB, EV, EM, BS, UE, PC
#
# I have no clue:
# * VI, MC, SS, IO, PT, MK, MT, ID, BI, LD, PL, PI, SI, CI, SX, PE, PS

class DvfSN(NamedTuple):
    name: str
    flg0: int
    flg1: int
    flg2: int
    address: int
    flg3: int
    flg4: int

    def parse(b: bytes) -> 'DvfSN':
        name, flg0, flg1, flg2, addr, nul1, flg3, flg4, nulls = struct.unpack(
            '<16sH2HHH2H10s', b)
        name = name.rstrip(b'\0').decode('utf-8')
        assert nul1 == 0
        assert all(n == 0 for n in nulls)

        return DvfSN(name, flg0, flg1, flg2, addr, flg3, flg4)

DvfES = DvfSN

class DvfVN(NamedTuple):
    name: str
    address: int
    int_type: int
    enable_reg_addr: int
    enable_bit_flag: int

    def parse(b: bytes) -> DvfVN:
        name, addr, ityp, enrg, enbt = struct.unpack('<16sHHHH', b)
        name = name.rstrip(b'\0').decode('utf-8')

        return DvfVN(name, addr, ityp, enrg, enbt)

class DvfMI(NamedTuple):
    ACCESS_R = 1
    ACCESS_W = 2
    ACCESS_X = 4
    ACCESS_RSV = 8 # reserved

    MREG_FLASH     = 0x00
    MREG_RAM       = 0x01
    # TODO: rama again? (0x2)
    # TODO: 0..ffffe, 0..fffff regions (0xb, 0xc)
    MREG_RESERVED  = 0x0e
    MREG_CPUREGS   = 0x0f
    MREG_SFR       = 0x10
    # TODO ffd0..ffdf (0x11)
    MREG_VECTOR    = 0x12
    MREG_CALLT     = 0x13
    MREG_CALLF     = 0x14
    MREG_SHORT     = 0x16 # "short direct addressing" in instruction set
    MREG_SHORT_SFR = 0x18 # above + is SFR
    MREG_SFR2      = 0x24
    MREG_ROM_ENTRY = 0x25
    MREG_MIRROR    = 0x26
    MREG_PRIVRAM   = 0x27
    MREG_ROM       = 0x28

    addr_start: int
    addr_end: int
    mreg_type: int
    access: int

    def parse(b: bytes) -> DvfMI:
        start, end, typ, acc, zeros = struct.unpack('<IIHHI', b)
        assert zeros == 0
        return DvfMI(start, end, typ, acc)

class DvfMD(NamedTuple):
    version: str # file format version?
    timestamp: str # also includes author

    def parse(b: bytes) -> DvfMD:
        assert len(b) >= 40
        ver, ts = b[:8], b[8:]
        ver = ver.rstrip(b'\0').decode('utf-8')
        ts  = ts.rstrip( b'\0').decode('utf-8')
        return DvfMD(ver, ts)

class DvfSC(NamedTuple):
    name: str
    addr1: int
    addr2: int

    def parse(b: bytes) -> DvfSC:
        name, a1, a2 = struct.unpack('<16sII', b)
        name = name.rstrip(b'\0').decode('utf-8')
        return DvfSC(name, a1, a2)

class DvfFL(NamedTuple):
    FL_TYPE_NBLOCKS = 1
    FL_TYPE_BLOCK   = 2

    fl_type: int
    index: int
    address: int
    extra: int

    def parse(b: bytes) -> DvfFL:
        typ, ind, addr, ex = struct.unpack('<HHHH', b)
        return DvfFL(typ, ind, addr, ex)

class DvfOD(NamedTuple):
    index: int
    subindex: int
    value: int

    def parse(b: bytes) -> DvfOD:
        ind, subind, val = struct.unpack('<HHI', b)
        return DvfOD(ind, subind, val)

class DvfSect(NamedTuple):
    SN = name2int("SN")
    VN = name2int("VN")
    MI = name2int("MI")
    ES = name2int("ES")
    MD = name2int("MD")

    SC = name2int("SC")
    FL = name2int("FL")
    OD = name2int("OD")

    def ENTSIZE(dvftype: int) -> Dict[int, int]:
        return {
            name2int(b'SN'):40,
            name2int(b'VN'):24,
            name2int(b'MI'):16,
            name2int(b'CG'):80,
            name2int(b'ES'):40, # 78K0R only?
            name2int(b'XD'):12, # 78K0R only?
            name2int(b'CT'):96,
            name2int(b'SC'):24,
            name2int(b'QB'):136 if (dvftype & 0xF0) == 0x10 else 108,
            name2int(b'EV'):8,
            name2int(b'FL'):8,
            name2int(b'OD'):8,
            name2int(b'AF'):16,
            name2int(b'SD'):16,
            name2int(b'EM'):12,
            name2int(b'MD'):48,
            name2int(b'BS'):16, # 78K0 only?
            name2int(b'UE'):12, # RL78 only?
            name2int(b'PC'):20, # idk
        }

    TYPES = {
        name2int('SN'): DvfSN,
        name2int('ES'): DvfES,
        name2int('VN'): DvfVN,
        name2int('MI'): DvfMI,
        name2int('MD'): DvfMD,
        name2int('SC'): DvfSC,
        name2int('FL'): DvfFL,
        name2int('OD'): DvfOD
    }

    name: int
    nentries: int
    offset: int

class Dvf(NamedTuple):
    MAGIC = 0x1a

    # TODO: add support for V850/RH850, RX
    MCU_78K0    =    0
    MCU_78K0R   = 0x10
    MCU_RL78x1y = 0x20
    MCU_RL78x2y = 0x30

    IDK_KNOWN_78K0 = {
        b'\0\x001\x02!\x02\0\0H\x11', # Kx2
        b'\0\x001\x02\0\x01\0\0H\x11', # Kx2-A, uPD78F8039
        b'\0\x001\x02\x01\x01\0\0H\x11', # Fx2
        b'\0\x001\x02\x01\x02\0\0H\x11', # Kx2-L
        b'\0\x001\x02\x10\x01\0\0H\x11', # misc (Ix2, KC2-A, uPD78F0730)
        b'\0\x001\x02\x11\x01\0\0H\x11', # FC2
        b'\0\x001\x02\x12\x01\0\0H\x11', # Lx3
        b'\0\x002\x02\0\x01\0\0H\x11', # misc (uPD78F8077, Dx2, Kx2-C, uPD78F8039)
        b'\0\x002\x02\x02\x01\0\0H\x11', # Fx2-L
        b'\0\x002\x02\x10\x01\0\0H\x11', # Lx3-M
    }
    IDK_KNOWN_78K0R = {
        b'\x0b\0\0\x03 \x02\0\0H\x11', # KC3-L
        b'\x0b\0\0\x03!\x01\0\0H\x11', # Lx3
        b'\x0b\0\0\x03\0\x01\0\0H\x11', # misc (Kx3-L(USB), Hx3, LG3-M, uPD78F80xx)
        b'\x0b\0\0\x03\0\x02\0\0H\x11', # Kx3-L, Ix3
        b'\x0b\0\0\x03\0\x03\0\0H\x11', # Kx3
        b'\x0b\0\0\x03\x01\x01\0\0H\x11', # Fx3, Kx3-C
        b'\x0b\0\0\x03\x01\x02\0\0H\x11', # Kx3-L
        b'\x0b\0\0\x03\x10\x01\0\0H\x11', # Lx3, Kx3-A
        b'\x0b\0\0\x03\x10\x02\0\0H\x11', # Kx3-L
        b'\x0b\0\0\x03\x11\x01\0\0H\x11', # KF3-C
    }
    IDK_KNOWN_RL78x1y = {
        b'\x0c\0\0\x03 \x01\0\0H\x11', # G11, G1D
        b'\x0c\0\0\x03!\x01\0\0H\x11', # I1A, L1C
        b'\x0c\0\0\x030\x02\0\0H\x11', # G14
        b'\x0c\0\0\x03\0\x01\0\0H\x11', # misc (G13A, G1H, F15, I1E, H1D, I1C)
        b'\x0c\0\0\x03\x01\x01\0\0H\x11', # misc (L1A, I1D, G1G, G1A, L13, I1C, D1A, G1F, F1A)
        b'\x0c\0\0\x03\x01\x02\0\0H\x11', # G1E
        b'\x0c\0\0\x03\x02\x01\0\0H\x11', # I1B, F12
        b'\x0c\0\0\x03\x03\x01\0\0H\x11', # G10, L12
        b'\x0c\0\0\x03\x10\x01\0\0H\x11', # misc (G1N, F1E, G1P, G1C, D1A, G1M, G1K)
        b'\x0c\0\0\x03\x11\x01\0\0H\x11', # L1C, F13, F14
        b'\x0c\0\0\x03\x12\x01\0\0H\x11', # G12
        b'\x0c\0\0\x03\x14\x01\0\0H\x11', # G13
    }
    IDK_KNOWN_RL78x2y = {
        b'\x0c\0 \x03\0\x01\0\0H\x11', # F24
        b'\x0c\0\0\x030\x02\0\0H\x11', # FGIC
        b'\x0c\0\0\x03\x02\x01\0\0H\x11', # FGIC
        b'\x0c\0\x10\x03\x10\x01\0\0H\x11', # G23
    }

    magic: int
    copyright_notice: str
    idk: bytes
    chipname_disp: str
    chipname_sym: str

    mcutype: int
    sections: Dict[int, Any]
    sections_raw: List[DvfSect]

    @staticmethod
    def parse(data: bytes, mcutype = None) -> Dvf:
        magic, crnot, idk, name_disp, name_sym, nsect = struct.unpack(
            '<B35s16s16s8sI', data[:80])
        assert magic == Dvf.MAGIC
        assert idk[14:16] in { b"\x02\0", b"\x03\0", b"\x04\0" }, str(idk)
        crnot = crnot.rstrip(b'\0').decode('utf-8')
        name_disp = name_disp.rstrip(b'\0').decode('utf-8')
        name_sym  = name_sym.rstrip(b'\0').decode('utf-8')

        #mcutype = None
        if mcutype is None:
            idkfix = idk[0:10]
            if idkfix in Dvf.IDK_KNOWN_78K0: mcutype = Dvf.MCU_78K0
            elif idkfix in Dvf.IDK_KNOWN_78K0R: mcutype = Dvf.MCU_78K0R
            elif idkfix in Dvf.IDK_KNOWN_RL78x1y: mcutype = Dvf.MCU_RL78x1y
            elif idkfix in Dvf.IDK_KNOWN_RL78x2y: mcutype = Dvf.MCU_RL78x2y
            else: assert False, ("unknown MCU type", idk)

        sects = {}
        sraw = []
        for i in range(nsect):
            shblob = data[(80+i*12):(80+(i+1)*12)]
            name, nentries, offset, zeros = struct.unpack('<HHH6s', shblob)
            if not all(z == 0 for z in zeros):
                continue # idk

            sraw.append(DvfSect(name, nentries, offset))
            if nentries == 0: continue

            #print("%c%c" % (chr(name&0xff),chr(name>>8)))
            entsize = DvfSect.ENTSIZE(mcutype)[name]

            sdblob = data[offset:(offset+entsize*nentries)]

            styp = DvfSect.TYPES.get(name)
            if styp is None: sects[name] = sdblob
            else:
                l = []
                for j in range(nentries):
                    sdiblob = sdblob[(j*entsize):((j+1)*entsize)]
                    #if name == name2int("MD"):
                    #    print("esz",entsize,"j",j,"blob",sdiblob)
                    l.append(DvfSect.TYPES[name].parse(sdiblob))
                sects[name] = l

        return Dvf(magic, crnot, idk, name_disp, name_sym, mcutype, sects, sraw)