"""Parser generico delle sezioni a marcatori dei file Creo (.prt/.asm). Reverse engineering.
Modello: ogni campo dichiarato per nome (e0 <tipo> nome\\0) riceve un id progressivo.
In modalita' posizionale il campo atteso e' next[ultimo id letto]; f7 <id> lo forza.
"""
import struct, sys, os


class Err(Exception):
    pass


class Node(list):
    """struct: lista di (nome, valore)"""

    def get(s, k, d=None):
        for n, v in s:
            if n == k:
                return v
        return d

    def all(s, k):
        return [v for n, v in s if n == k]


class Arr(list):
    pass


class Unk(float):
    """costante a un byte non ancora identificata"""
    def __new__(cls, code):
        o = float.__new__(cls, 'nan')
        o.code = code
        return o

    def __repr__(s):
        return f"<c{s.code:02x}>"


class P:
    def __init__(s, b, trace=False):
        s.b = b
        s.p = 0
        s.fields = {}
        s.nid = 0
        s.next = {}
        s.alias = {}
        s.stack = []
        s.trace = trace
        s.root = None
        s.last = None
        s.lenient = 0
        s.lastc = None
        s.popped = None

    def u8(s):
        v = s.b[s.p]
        s.p += 1
        return v

    def peek(s):
        return s.b[s.p] if s.p < len(s.b) else None

    def rint(s):
        b = s.u8()
        if b < 0x80:
            return b
        if b < 0xc0:
            return (b & 0x3f) << 8 | s.u8()
        if b < 0xe0:
            v = (b & 0x1f) << 16 | s.b[s.p] << 8 | s.b[s.p + 1]
            s.p += 2
            return v
        if b == 0xe8:
            return s.u8()
        if b == 0xe9:
            v = int.from_bytes(s.b[s.p:s.p + 2], 'big')
            s.p += 2
            return v
        if b == 0xea:
            v = int.from_bytes(s.b[s.p:s.p + 3], 'big')
            s.p += 3
            return v
        if b == 0xeb:
            v = int.from_bytes(s.b[s.p:s.p + 4], 'big')
            s.p += 4
            return v
        if b == 0xf5:
            return s.rint() | 0x80000000
        if b == 0xf6:
            return -1
        raise Err(f"int tag {b:02x} @{s.p - 1:x}")

    CONST = {0x17: -6.283185307179586, 0x14: 4.71238898038469, 0x0d: -0.5, 0x11: 3.141592653589793, 0x15: -3.141592653589793, 0x12: 1.5707963267948966, 0x16: -1.5707963267948966, 0x13: 6.283185307179586,  0x18: 0.0, 0x0f: 1.0, 0x0e: 0.5, 0x10: -1.0, 0x0c: 0.1, 0x0a: 0.01, 0x07: 0.001}

    FBASE = 0x3f9d
    FCONST = {0x0b: 1.0, 0x28: -1.0}

    def rflt(s):
        t = s.u8()
        if 0x30 <= t <= 0x4f:
            T = (0x37 + t - 0x30) if t < 0x40 else (0xb7 + t - 0x40)
            raw = bytes([T]) + s.b[s.p:s.p + 3]
            s.p += 3
        elif t == 0xea:
            raw = s.b[s.p:s.p + 3] + bytes(1)
            s.p += 3
        elif t == 0xeb:
            raw = s.b[s.p:s.p + 4]
            s.p += 4
        elif t == 0xe9:
            raw = s.b[s.p:s.p + 2] + bytes(2)
            s.p += 2
        elif t == 0xe8:
            raw = s.b[s.p:s.p + 1] + bytes(3)
            s.p += 1
        elif t == 0x00:
            return 0.0
        elif t in s.FCONST:
            return s.FCONST[t]
        elif t < 0x30:
            return Unk(t)
        elif 0x50 <= t <= 0xdf:
            raw = (s.FBASE + t).to_bytes(2, 'big') + s.b[s.p:s.p + 2]
            s.p += 2
        else:
            raise Err(f"flt tag {t:02x} @{s.p - 1:x}")
        return struct.unpack('>f', raw)[0]

    def rbyte(s):
        b = s.u8()
        if b == 0xe8:
            return s.u8()
        return b

    # prefissi negativi a 16 bit dei tag a5-df: bfd0..c00d meno tre buchi (bfd2, bfdf, bffe)
    NEG = [x for x in range(0xbfd0, 0xc00e) if x not in (0xbfd2, 0xbfdf, 0xbffe)]

    def rdbl(s):
        t = s.u8()
        if t in s.CONST:
            return s.CONST[t]
        if t <= 0x16 and t > 0x03:
            return Unk(t)
        if t <= 0x03:
            c = s.u8()
            k = t + 1
            pl = s.b[s.p:s.p + k]
            s.p += k
            if c not in s.CONST:
                if 0x04 <= c <= 0x16:
                    return Unk(c)
                raise Err(f"dbl const {c:02x} after modifier @{s.p - 1:x}")
            raw = bytearray(struct.pack('>d', s.CONST[c]))
            raw[8 - k:] = pl
            return struct.unpack('>d', bytes(raw))[0]
        if t == 0xed:
            v = struct.unpack('>d', s.b[s.p:s.p + 8])[0]
            s.p += 8
            return v
        if t == 0xea:
            v = struct.unpack('>d', s.b[s.p:s.p + 3] + bytes(5))[0]
            s.p += 3
            return v
        if t == 0xec:
            v = struct.unpack('>d', s.b[s.p:s.p + 7] + bytes(1))[0]
            s.p += 7
            return v
        if 0x19 <= t <= 0x48:
            if t <= 0x2f:
                T = 0x3c + (t - 0x17) // 5
                m = (t - 0x17) % 5
            else:
                T = 0xbc + (t - 0x30) // 5
                m = (t - 0x30) % 5
            if m == 4:
                pl = s.b[s.p:s.p + 2]
                s.p += 2
                raw = bytes([T]) + pl + b'\x00' * 5
            elif m == 3:
                pl = s.b[s.p:s.p + 2]
                s.p += 2
                raw = bytes([T]) + pl + pl[1:] * 5
            elif m == 2:
                pl = s.b[s.p:s.p + 7]
                s.p += 7
                raw = bytes([T]) + pl
            elif m == 1:
                pl = s.b[s.p:s.p + 6]
                s.p += 6
                raw = bytes([T]) + pl + b'\x00'
            else:
                pl = s.b[s.p:s.p + 3]
                s.p += 3
                raw = bytes([T]) + pl + bytes(4)
            return struct.unpack('>d', raw)[0]
        if 0x49 <= t <= 0xdf:
            hi = (0x3f75 + t) if t <= 0xa4 else s.NEG[t - 0xa5]
            pl = s.b[s.p:s.p + 6]
            s.p += 6
            return struct.unpack('>d', hi.to_bytes(2, 'big') + pl)[0]
        raise Err(f"dbl tag {t:02x} @{s.p - 1:x}")

    def rstr(s):
        if s.peek() == 0xe1:
            s.p += 1
            return None
        e = s.b.index(0, s.p)
        v = s.b[s.p:e].decode('latin1')
        s.p = e + 1
        return v

    def rarr(s, n, rd):
        out = []
        while len(out) < n:
            b = s.peek()
            if 0xe4 <= b <= 0xe6 and out:
                s.p += 1
                out += [out[-1]] * (b - 0xe3)
            elif b == 0xe7 and out:
                s.p += 1
                k = s.rint()
                out += [out[-1]] * k
            else:
                out.append(rd())
        if len(out) != n:
            raise Err(f"array overrun {len(out)}>{n} @{s.p:x}")
        return out

    def scalar(s, f):
        t = f['t']
        if t in (0x0a, 0x03):
            if s.peek() == 0xf8:
                s.p += 1
                f['dim'] = (s.rint(),)
                return [s.rstr() for _ in range(f['dim'][0])]
            if f.get('dim'):
                out = []
                while len(out) < f['dim'][0] and 0x20 <= s.peek() < 0xe0:
                    out.append(s.rstr())
                return out
            return s.rstr()
        rd = s.rdbl if t == 0x02 else s.rbyte if t in (0x07, 0x04) else s.rflt if t == 0x06 else s.rint
        b = s.peek()
        if b == 0xe1:
            s.p += 1
            return None
        if b == 0xf8:
            s.p += 1
            f['dim'] = (s.rint(),)
        elif b == 0xf9:
            s.p += 1
            f['dim'] = (s.rint(), s.rint())
        elif b == 0xfc:
            s.p += 1
            f['dim'] = (s.rint(),) + tuple(f['dim'][1:])
        elif b == 0xfa:
            s.p += 1
            f['dim'] = tuple(s.rint() for _ in range(s.rint()))
        d = f.get('dim')
        if d is None:
            return rd()
        n = 1
        for k in d:
            n *= k
        v = s.rarr(n, rd)
        if len(d) == 2:
            v = [v[i * d[1]:(i + 1) * d[1]] for i in range(d[0])]
        return v

    def sch(s, fid):
        return s.alias.get(fid, fid)

    # --- contenitori -------------------------------------------------
    def owner_index(s, E, post=False):
        own = s.fields[E]['own']
        for i in range(len(s.stack) - 1, -1, -1):
            c = s.stack[i]
            if c['kind'] == 'struct':
                if c['schema'] in own:
                    if post and s.fields[E].get('post') and i > 0:
                        return i - 1
                    return i
            elif c['occ'] == E or ('A', c['occ']) in own:
                return i
        return None

    def default_close(s):
        """chiusura implicita prima di un campo con nome o di un salto senza marcatore f0-f4:
        se il campo atteso appartiene al contenitore padre di quello in cui e' stato scritto
        l'ultimo campo, quest'ultimo contenitore si considera chiuso."""
        if s.last is None or not s.stack or s.lastc is None:
            return
        E = s.next.get(s.last)
        if E is None:
            return
        p = None
        for k in range(len(s.stack) - 1, -1, -1):
            if s.stack[k] is s.lastc:
                p = k
                break
        if not p:
            return
        i = s.owner_index(E, True)
        if i == p - 1:
            del s.stack[p:]

    def marker(s):
        b = s.peek()
        if b is None or not (0xf0 <= b <= 0xf4):
            return None
        s.p += 1
        if b == 0xf4:
            return s.rint()
        return b - 0xf0

    def explicit(s, n):
        for _ in range(n):
            if not s.stack:
                raise Err(f"pop underflow @{s.p:x}")
            s.popped = s.stack.pop()

    def trans(s, E):
        if s.last is not None:
            s.next[s.last] = E
        s.last = E

    def locate(s, E, jump, stay):
        """porta in cima allo stack il contenitore che riceve il campo E. Ritorna 'elem' o 'field'."""
        f = s.fields[E]
        hdr = s.peek() in (0xf8, 0xf9)
        top = s.stack[-1]
        if not hdr and f['t'] == 0:
            for i in range(len(s.stack) - 1, -1, -1):
                c = s.stack[i]
                if c['kind'] == 'array' and c['occ'] == E:
                    del s.stack[i + 1:]
                    return 'elem'
        par = f['parent']
        i = s.owner_index(E, not jump)
        if i is None and jump and top['kind'] == 'struct' and top['node']:
            # variante di unione riusata in un altro contenitore: resta dove siamo
            f['own'].add(top['schema'])
            s.lenient += 1
            return 'field'
        fresh = top['kind'] == 'struct' and not top['node']
        if jump and fresh and i != len(s.stack) - 1:
            s.alias[top['field']] = par
            top['schema'] = par
            return 'field'
        if i is not None:
            if jump and i != len(s.stack) - 1:
                s.lenient += 1
            del s.stack[i + 1:]
            return 'field'
        if jump and top['kind'] == 'array' and not top['node'] and f['t'] == 0 and not hdr:
            top['occ'] = E
            return 'elem'
        raise Err(f"field {E} {f['name']}: container {par} not on stack @{s.p:x}")

    def element(s, E):
        c = s.stack[-1]
        s.trans(E)
        s.lastc = c
        if s.peek() == 0xfb:
            s.p += 1
        b = s.u8()
        if len(c['node']) >= c['n']:
            raise Err(f"too many elements in {s.fields[c['field']]['name']} @{s.p:x}")
        if b == 0xe1:
            c['node'].append(None)
        elif b in (0xe2, 0xe3):
            v = Node()
            c['node'].append(v)
            s.stack.append(dict(kind='struct', field=c['field'], schema=s.sch(c['occ']), node=v, elem=True, occ=None, occ0=c['occ']))
        else:
            raise Err(f"array elem tag {b:02x} @{s.p - 1:x} in {s.fields[c['field']]['name']}")

    def value(s, fid):
        c = s.stack[-1] if s.stack else None
        f = s.fields[fid]
        s.trans(fid)
        s.lastc = c
        if f['t'] != 0:
            v = s.scalar(f)
        else:
            b = s.u8()
            if b == 0xe1:
                v = None
            elif b in (0xe2, 0xe3):
                v = Node()
            elif b == 0xf8:
                v = Arr()
                n = s.rint()
            elif b == 0xf9:
                v = Arr()
                n = s.rint()
                s.rint()
            else:
                raise Err(f"ptr tag {b:02x} @{s.p - 1:x} field {f['name']}")
        if c is None:
            s.root = (f['name'], v)
        elif c['kind'] == 'array':
            c.setdefault('tail', []).append((f['name'], v))
        else:
            c['node'].append((f['name'], v))
        if isinstance(v, Arr):
            s.stack.append(dict(kind='array', field=fid, occ=fid, n=n, node=v))
        elif isinstance(v, Node):
            s.stack.append(dict(kind='struct', field=fid, schema=s.sch(fid), node=v, occ=None))

    def run(s):
        b = s.b
        while s.p < len(b):
            m = s.marker()
            x = s.peek()
            if x is None:
                s.explicit(m or 0)
                break
            if x == 0xe0:
                s.p += 1
                t = s.u8()
                e = b.index(0, s.p)
                name = b[s.p:e].decode('latin1')
                s.p = e + 1
                m2 = s.marker()
                if m2 is not None:
                    m = (m or 0) + m2
                s.popped = None
                if m is None:
                    s.default_close()
                else:
                    s.explicit(m)
                c = s.stack[-1] if s.stack else None
                s.nid += 1
                fid = s.nid
                if c is None:
                    par = 0
                elif c['kind'] == 'array':
                    par = ('A', c['occ'])
                else:
                    par = c['schema']
                s.fields[fid] = dict(name=name, t=t, parent=par, own={par})
                pp = s.popped
                if pp is not None and c is not None and pp['kind'] == 'struct' and c['kind'] == 'struct'                         and pp['schema'] == c['schema'] and t == 0:
                    s.fields[fid]['post'] = True
                if s.trace:
                    print(f"{s.p:6x} {'  ' * len(s.stack)}#{fid} {name} t={t}")
                s.value(fid)
                continue
            if m:
                s.explicit(m)
            jump = False
            if x == 0xf7:
                s.p += 1
                E = s.rint()
                jump = True
                if m is None:
                    s.default_close()
                if E not in s.fields:
                    raise Err(f"jump to unknown field {E} @{s.p:x}")
            else:
                E = s.next.get(s.last)
                if E is None:
                    raise Err(f"no expected field @{s.p:x} after {s.last} "
                              f"{s.fields[s.last]['name'] if s.last else ''}")
            kind = s.locate(E, jump or m is not None, m == 0)
            if s.trace:
                print(f"{s.p:6x} {'  ' * len(s.stack)}{'>' if jump else '='}{E} {s.fields[E]['name']} {kind}")
            if kind == 'elem':
                s.element(E)
            else:
                s.value(E)
        return s.root


def short(x, maxarr=40):
    if isinstance(x, Unk):
        return repr(x)
    if isinstance(x, float):
        return repr(round(x, 9))
    if isinstance(x, list):
        r = ', '.join(short(y) for y in x[:maxarr])
        return f"[{r}{', ...' if len(x) > maxarr else ''}] (n={len(x)})"
    return repr(x)


def fmt(v, ind=0, out=None, maxarr=40):
    pad = '  ' * ind
    if isinstance(v, Node):
        for n, x in v:
            if isinstance(x, (Node, Arr)):
                out.append(f"{pad}{n}:")
                fmt(x, ind + 1, out, maxarr)
            else:
                out.append(f"{pad}{n} = {short(x, maxarr)}")
    elif isinstance(v, Arr):
        for i, x in enumerate(v):
            out.append(f"{pad}[{i}]")
            fmt(x, ind + 1, out, maxarr)


def parse(b, trace=False):
    p = P(b, trace)
    try:
        root = p.run()
        return p, root, None
    except (Err, IndexError, ValueError, KeyError) as e:
        return p, p.root, e


if __name__ == '__main__':
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    import prt
    d, secs, order = prt.load(sys.argv[1])
    names = sys.argv[2].split(',') if len(sys.argv) > 2 and sys.argv[2] != '*' else [n for n in order if secs[n][2][2] == '9']
    outdir = sys.argv[3] if len(sys.argv) > 3 and sys.argv[3] != '-' else None
    for nme in names:
        o, l, f = secs[nme]
        b = d[o:o + l]
        p, root, e = parse(b, trace=len(sys.argv) > 4)
        ok = e is None and not p.stack
        print(f"{nme:22} len={l:7} parsed={p.p:7} fields={p.nid:5} lenient={p.lenient} {'OK' if ok else 'ERR ' + repr(e) + ' stack=' + str(len(p.stack))}")
        if e is not None:
            prt.hexdump(b, max(0, p.p - 64), 128)
        if outdir and root:
            out = [root[0] + ':']
            fmt(root[1], 1, out)
            open(f"{outdir}/{nme.replace(':', '_').replace('#', '_')}.txt", 'w', encoding='utf-8').write('\n'.join(out))