#!/usr/bin/env python3
"""
b2z.py -- reader for DIVISION / Virtual World Entertainment "DIV-BIZ2" (.B2Z / .V2Z)
binary geometry objects, used by the DPL3 renderer (c. 1994).

Faithful transcription of DPL3/BIZREAD.C (Copyright DIVISION Ltd 1994, author PJA).
The original read the file in 8 KB blocks with endian correction; here we load the
whole file into memory and walk an absolute cursor -- identical results, less code.

Little-endian throughout (the files were authored on x86 / i860 hosts).

Usage:
    python b2z.py dump  <file.b2z>            # structural dump
    python b2z.py obj    <file.b2z> [out.obj] # export triangulated Wavefront OBJ
"""
import struct
import sys
from dataclasses import dataclass, field


# ----------------------------------------------------------------------------
# low-level stream (mirrors get_char / get_int* / get_float32 / get_struct32)
# ----------------------------------------------------------------------------
class Stream:
    def __init__(self, data: bytes):
        self.d = data
        self.pos = 0

    def eof(self) -> bool:
        return self.pos >= len(self.d)

    def u8(self) -> int:
        v = self.d[self.pos]
        self.pos += 1
        return v

    def i8(self) -> int:
        v = self.u8()
        return v - 256 if v >= 128 else v

    def u16(self) -> int:
        v = struct.unpack_from("<H", self.d, self.pos)[0]
        self.pos += 2
        return v

    def i32(self) -> int:
        v = struct.unpack_from("<i", self.d, self.pos)[0]
        self.pos += 4
        return v

    def f32(self) -> float:
        v = struct.unpack_from("<f", self.d, self.pos)[0]
        self.pos += 4
        return v

    def floats(self, n: int):
        v = struct.unpack_from("<%df" % n, self.d, self.pos)
        self.pos += 4 * n
        return list(v)

    def cstr(self) -> str:
        """null-terminated string (the `while (c=get_char())` idiom)."""
        start = self.pos
        while self.d[self.pos] != 0:
            self.pos += 1
        s = self.d[start:self.pos].decode("latin-1")
        self.pos += 1  # skip the NUL
        return s

    def fixed_str(self, n: int) -> str:
        s = self.d[self.pos:self.pos + n].decode("latin-1")
        self.pos += n
        return s

    def skip(self, n: int):
        self.pos += n


# ----------------------------------------------------------------------------
# block header: 16-bit type; top nibble selects the width of the length field.
#   (hdr >> 12) & 0xc  ==  0x8 -> int32 len,  0x4 -> int16 len,  0x0 -> int8 len
# The record TYPE is (hdr & 0xfff); the top nibble is only a length-encoding flag,
# which is why every record type appears as 0x0NNN / 0x4NNN / 0x8NNN in the source.
# ----------------------------------------------------------------------------
def read_block(s: Stream):
    hdr = s.u16()
    sel = (hdr >> 12) & 0xc
    if sel == 0x8:
        length = s.i32() & 0xffffffff
    elif sel == 0x4:
        length = s.u16()
    elif sel == 0x0:
        length = s.u8()
    else:
        raise ValueError("Unrecognised block header 0x%04x" % hdr)
    return hdr, length


# ----------------------------------------------------------------------------
# data model (subset of dpltypes.h that a .b2z can express)
# ----------------------------------------------------------------------------
@dataclass
class Vertex:
    pos: tuple            # (x, y, z)
    normal: tuple = None  # (nx, ny, nz)
    rgba: tuple = None    # (r, g, b, a)
    lum: tuple = None     # (l, a)
    uv: tuple = None      # (u, v) or (u, v, w)


# geometry primitive types (dpl_geo_type)
TRISTRIP, POLYSTRIP, PMESH, POLYGON = "tristrip", "polystrip", "pmesh", "polygon"


@dataclass
class Geometry:
    geotype: str
    verts: list = field(default_factory=list)
    conns: list = field(default_factory=list)  # explicit faces (pmesh); each = list[int]


@dataclass
class Geogroup:
    name: str = ""
    draw_mode: int = 0
    f_material: str = None
    b_material: str = None
    geoms: list = field(default_factory=list)


@dataclass
class Material:
    name: str
    ambient: tuple = (1, 1, 1)
    diffuse: tuple = (1, 1, 1)
    specular: tuple = (0, 0, 0, 0)
    emissive: tuple = (1, 1, 1)
    opacity: tuple = (1, 1, 1)
    texture: str = None
    ramp: str = None


@dataclass
class Texture:
    name: str
    mapfile: str = None
    minify: int = 0
    magnify: int = 0
    alpha: int = 0
    wrap_u: int = 0
    wrap_v: int = 0
    bitslice: int = 0   # tag 0x018 (1996+): 4-bit plane index into a .BSL pack
    special: str = None  # tag 0x037: free-text hook, e.g. " SCROLL u0 v0 du dv"


@dataclass
class Model:
    object_name: str = ""
    scale: float = 1.0
    units: float = 1.0
    geogroups: list = field(default_factory=list)
    materials: dict = field(default_factory=dict)
    textures: dict = field(default_factory=dict)
    ramps: dict = field(default_factory=dict)


# vertex-format table: header type (12-bit) -> (floats_per_vertex, offsets dict)
# offsets are indices into the per-vertex float array. Mirrors parse_vertices().
VFMT = {
    0x080: (3,  {}),                               # flat  x y z
    0x081: (6,  {"n": 3}),                          # + normal
    0x082: (7,  {"c": 3}),                          # + rgba (cooked)
    0x083: (10, {"n": 3, "c": 6}),                  # normal + rgba
    0x084: (5,  {"l": 3}),                          # + luminance/alpha
    0x085: (8,  {"n": 3, "l": 6}),                  # normal + la
    0x088: (5,  {"t2": 3}),                         # + 2d texture uv
    0x089: (8,  {"n": 3, "t2": 6}),                 # normal + uv
    0x08a: (9,  {"c": 3, "t2": 7}),                 # rgba + uv
    0x08c: (7,  {"l": 3, "t2": 5}),                 # la + uv
    0x090: (6,  {"t3": 3}),                         # + 3d texture uvw
    0x091: (9,  {"n": 3, "t3": 6}),                 # normal + uvw
    0x092: (10, {"c": 3, "t3": 7}),                 # rgba + uvw
    0x094: (8,  {"l": 3, "t3": 5}),                 # la + uvw
}

PRIM = {0x043: POLYGON, 0x044: TRISTRIP, 0x045: POLYSTRIP, 0x046: PMESH}


class Reader:
    def __init__(self, data: bytes):
        self.s = Stream(data)
        self.m = Model()

    # -- header ------------------------------------------------------------
    def parse_header(self, length):
        end = self.s.pos + length
        while self.s.pos < end:
            hdr, ln = read_block(self.s)
            body = self.s.pos
            t = hdr & 0xfff
            if t == 0x005:      # scale (0x2005)
                self.m.scale = self.s.f32()
            elif t == 0x006:    # precision (0x2006): 0 = single precision
                if self.s.u8() != 0:
                    raise ValueError("double-precision files unsupported (as in original)")
            elif t == 0x009:    # unit (0x2009)
                u = self.s.u8()
                self.m.units = 1.0 if u == 0 else (1.0 / 25.4)
            # version/date/time/filetype -> skipped
            self.s.pos = body + ln

    # -- vertices ----------------------------------------------------------
    def read_vertices(self, fpv, offs, n):
        out = []
        for _ in range(n):
            f = self.s.floats(fpv)
            v = Vertex(pos=(f[0], f[1], f[2]))
            if "n" in offs:
                o = offs["n"]; v.normal = (f[o], f[o+1], f[o+2])
            if "c" in offs:
                o = offs["c"]; v.rgba = (f[o], f[o+1], f[o+2], f[o+3])
            if "l" in offs:
                o = offs["l"]; v.lum = (f[o], f[o+1])
            if "t2" in offs:
                o = offs["t2"]; v.uv = (f[o], f[o+1])
            if "t3" in offs:
                o = offs["t3"]; v.uv = (f[o], f[o+1], f[o+2])
            out.append(v)
        return out

    def parse_vertices(self, length, geotype, group):
        end = self.s.pos + length
        geom = Geometry(geotype)
        pre_conns = []          # pmesh connectivity accumulates across sub-blocks
        while self.s.pos < end:
            hdr, ln = read_block(self.s)
            body = self.s.pos
            t = hdr & 0xfff
            nfloats = ln >> 2
            if t == 0x047:                       # pmesh triangle connectivity
                for _ in range(ln // 12):
                    pre_conns.append([self.s.i32(), self.s.i32(), self.s.i32()])
            elif t == 0x04d:                     # pmesh polygon connectivity
                vpp = self.s.u8()
                for _ in range((ln - 1) // (vpp * 4)):
                    pre_conns.append([self.s.i32() for _ in range(vpp)])
            elif t in VFMT:                      # a vertex block
                fpv, offs = VFMT[t]
                geom.verts = self.read_vertices(fpv, offs, nfloats // fpv)
            # sphere/line/text (0x048/0x04a/0x04b) not implemented in original
            self.s.pos = body + ln
        if geotype == PMESH:
            geom.conns = pre_conns
        group.geoms.append(geom)

    # -- patch (geogroup) --------------------------------------------------
    def parse_patch(self, length, draw_mode):
        end = self.s.pos + length
        g = Geogroup(draw_mode=draw_mode)
        while self.s.pos < end:
            hdr, ln = read_block(self.s)
            body = self.s.pos
            t = hdr & 0xfff
            if t == 0x008:                       # name (0x2008)
                g.name = self.s.cstr()
            elif t == 0x030:                     # front material (0x2030)
                g.f_material = self._material_ref()
            elif t == 0x031:                     # back material (0x2031)
                g.b_material = self._material_ref(back=True, front=g.f_material)
            elif t == 0x036:                     # vertex format / draw mode
                g.draw_mode = self.s.u8()
            elif t in PRIM:                      # a geometry primitive
                self.s.pos = body            # parse_vertices re-reads from body
                # length already known (ln); rewind not needed -- body is content start
                self.parse_vertices(ln, PRIM[t], g)
            # plane/decal/facet/voodoo/sphere/line/text -> skipped
            self.s.pos = body + ln
        self.m.geogroups.append(g)

    def _material_ref(self, back=False, front=None):
        typ = self.s.u8()
        if typ == 0:
            return None
        if typ == 2:
            return "DEFAULT"
        if typ == 3 and back:
            return front
        if typ == 1:
            return self.s.cstr()
        return None

    # -- object ------------------------------------------------------------
    def parse_object(self, length):
        self._lod_taken = False
        self._object_body(self.s.pos + length, 0)

    def _object_body(self, end, draw_mode, in_lod=False):
        while self.s.pos < end:
            hdr, ln = read_block(self.s)
            body = self.s.pos
            t = hdr & 0xfff
            if t == 0x008:                       # object name
                self.m.object_name = self.s.cstr()
            elif t == 0x036:                     # vertex format -> draw_mode
                draw_mode = self.s.u8()
            elif t == 0x042:                     # patch
                self.s.pos = body
                self.parse_patch(ln, draw_mode)
            elif t == 0x041 and not in_lod:
                # LOD. In the 1994 format its payload was empty (a marker the
                # original loader skipped); the 1996 game format nests the
                # patches INSIDE it, with 0x046(len 8) = switch in/out floats.
                # Parse only the FIRST (highest-detail) LOD -- later ones are
                # lower-poly alternates and would double-draw.
                if not self._lod_taken:
                    self._lod_taken = True
                    self._object_body(body + ln, draw_mode, in_lod=True)
            # materials-on-object / comments -> skipped
            self.s.pos = body + ln

    # -- texture / material / ramp ----------------------------------------
    def parse_texture(self, length):
        end = self.s.pos + length
        tx = Texture(name="")
        while self.s.pos < end:
            hdr, ln = read_block(self.s)
            body = self.s.pos
            t = hdr & 0xfff
            if t == 0x008:
                tx.name = self.s.cstr()
            elif t == 0x011:
                tx.mapfile = self.s.cstr()
            elif t == 0x012:
                tx.minify = self.s.u8()
            elif t == 0x013:
                tx.magnify = self.s.u8()
            elif t == 0x014:
                tx.alpha = self.s.u8()
            elif t == 0x015:
                tx.wrap_u = self.s.u8()
            elif t == 0x016:
                tx.wrap_v = self.s.u8()
            elif t == 0x018:
                tx.bitslice = self.s.u8()
            elif t == 0x037:
                tx.special = self.s.cstr()
            self.s.pos = body + ln
        self.m.textures[tx.name] = tx

    def parse_material(self, length):
        end = self.s.pos + length
        mt = Material(name="")
        while self.s.pos < end:
            hdr, ln = read_block(self.s)
            body = self.s.pos
            t = hdr & 0xfff
            if t == 0x008:
                mt.name = self.s.cstr()
            elif t == 0x021:                     # texture ref
                mode = self.s.u8()
                if mode == 2:
                    mt.texture = self.s.cstr()
            elif t == 0x023:
                mt.ambient = tuple(self.s.floats(3))
            elif t == 0x024:
                mt.diffuse = tuple(self.s.floats(3))
            elif t == 0x025:
                mt.specular = tuple(self.s.floats(4))
            elif t == 0x026:
                mt.emissive = tuple(self.s.floats(3))
            elif t == 0x027:
                mt.opacity = tuple(self.s.floats(3))
            elif t == 0x028:
                mt.ramp = self.s.cstr()
            self.s.pos = body + ln
        self.m.materials[mt.name] = mt

    def parse_ramp(self, length):
        end = self.s.pos + length
        name, c0, c1 = "", (0, 0, 0), (0, 0, 0)
        while self.s.pos < end:
            hdr, ln = read_block(self.s)
            body = self.s.pos
            t = hdr & 0xfff
            if t == 0x008:
                name = self.s.cstr()
            elif t == 0x031:
                c0 = tuple(self.s.floats(3))
                c1 = tuple(self.s.floats(3))
            self.s.pos = body + ln
        self.m.ramps[name] = (c0, c1)

    # -- top level ---------------------------------------------------------
    def read(self):
        magic = self.s.fixed_str(8)
        if magic != "DIV-BIZ2":
            raise ValueError("Not a DIV-BIZ2 file (magic=%r)" % magic)
        while not self.s.eof():
            hdr, length = read_block(self.s)
            t = hdr & 0xfff
            if t == 0x005:       # trailer
                break
            elif t == 0x003:     # header
                self.parse_header(length)
            elif t == 0x010:     # texture
                self.parse_texture(length)
            elif t == 0x020:     # material
                self.parse_material(length)
            elif t == 0x030:     # ramp
                self.parse_ramp(length)
            elif t == 0x040:     # object
                self.parse_object(length)
            elif t == 0x004:     # comment
                self.s.skip(length)
            else:
                self.s.skip(length)
        return self.m


def load(path: str) -> Model:
    with open(path, "rb") as fp:
        return Reader(fp.read()).read()


# ----------------------------------------------------------------------------
# triangulation for export
# ----------------------------------------------------------------------------
def triangles(geom: Geometry):
    """Yield (i0, i1, i2) index triples in the geometry's own vertex space."""
    n = len(geom.verts)
    if geom.geotype == TRISTRIP:
        for i in range(n - 2):
            yield (i, i + 1, i + 2) if i % 2 == 0 else (i + 1, i, i + 2)
    elif geom.geotype == PMESH:
        for face in geom.conns:
            for k in range(1, len(face) - 1):
                yield (face[0], face[k], face[k + 1])
    elif geom.geotype in (POLYGON, POLYSTRIP):
        # a strip of polygons stored as a fan of the block's vertices
        for k in range(1, n - 1):
            yield (0, k, k + 1)


def export_obj(model: Model, out) -> tuple:
    base = 1  # OBJ indices are 1-based
    nv = nf = 0
    out.write("# exported from DIV-BIZ2 by b2z.py\n")
    out.write("# object: %s\n" % model.object_name)
    for gi, g in enumerate(model.geogroups):
        out.write("g %s\n" % (g.name or ("geogroup_%d" % gi)))
        for geom in g.geoms:
            for v in geom.verts:
                out.write("v %.6g %.6g %.6g\n" % v.pos)
            has_uv = any(v.uv for v in geom.verts)
            if has_uv:
                for v in geom.verts:
                    u, w = (v.uv[0], v.uv[1]) if v.uv else (0.0, 0.0)
                    out.write("vt %.6g %.6g\n" % (u, w))
            for (a, b, c) in triangles(geom):
                if has_uv:
                    out.write("f %d/%d %d/%d %d/%d\n" %
                              (base+a, base+a, base+b, base+b, base+c, base+c))
                else:
                    out.write("f %d %d %d\n" % (base+a, base+b, base+c))
                nf += 1
            base += len(geom.verts)
            nv += len(geom.verts)
    return nv, nf


# ----------------------------------------------------------------------------
# CLI
# ----------------------------------------------------------------------------
def cmd_dump(path):
    m = load(path)
    print("file       :", path)
    print("object     :", m.object_name or "(unnamed)")
    print("scale/units: %g / %g" % (m.scale, m.units))
    print("materials  : %d  %s" % (len(m.materials), list(m.materials)))
    print("textures   : %d  %s" % (len(m.textures), list(m.textures)))
    print("ramps      : %d  %s" % (len(m.ramps), list(m.ramps)))
    print("geogroups  : %d" % len(m.geogroups))
    tot_v = tot_t = 0
    for gi, g in enumerate(m.geogroups):
        print("  [%d] name=%r draw_mode=0x%x f_mtl=%r b_mtl=%r geoms=%d"
              % (gi, g.name, g.draw_mode, g.f_material, g.b_material, len(g.geoms)))
        for geom in g.geoms:
            tris = list(triangles(geom))
            tot_v += len(geom.verts); tot_t += len(tris)
            attrs = []
            if geom.verts:
                v0 = geom.verts[0]
                for a in ("normal", "rgba", "lum", "uv"):
                    if getattr(v0, a) is not None:
                        attrs.append(a)
            print("       - %-9s verts=%-4d faces=%-4d tris=%-4d attrs=%s"
                  % (geom.geotype, len(geom.verts), len(geom.conns), len(tris),
                     ",".join(attrs) or "pos"))
    print("TOTAL      : %d verts, %d triangles" % (tot_v, tot_t))
    for name, mt in m.materials.items():
        print("  material %-16s diffuse=%s tex=%s" %
              (name, tuple(round(x, 3) for x in mt.diffuse), mt.texture))


def cmd_obj(path, outpath=None):
    m = load(path)
    outpath = outpath or (path.rsplit(".", 1)[0] + ".obj")
    with open(outpath, "w") as fp:
        nv, nf = export_obj(m, fp)
    print("wrote %s : %d vertices, %d triangles" % (outpath, nv, nf))


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print(__doc__)
        sys.exit(1)
    cmd, path = sys.argv[1], sys.argv[2]
    if cmd == "dump":
        cmd_dump(path)
    elif cmd == "obj":
        cmd_obj(path, sys.argv[3] if len(sys.argv) > 3 else None)
    else:
        print("unknown command", cmd); sys.exit(1)