#!/usr/bin/env python3
"""OBJ -> new world geometry (TRAK + tile part of MAP) in a copy of a level WAD.

  geom_import.py base.wad mesh.obj new.wad [--default-tex N] [--no-coll] [--mat N]
                 [--keep-tiles] [--band 6]

Input OBJ in game space as geom_export.py writes it (right-handed, y up, 1 unit = 1 tile, z > 0
lies in the grid; grid cell (cx,cz) = (floor(x), floor(z)), must be within 0..W-1 / 0..H-1), i.e.
as Blender shows it. --export-space reads the old convention (raw export space, z < 0).

Textures: the engine has NO free UVs. Each triangle shows one half of a texture
rectangle (entry of the TEXT table); the corners are chosen via face flags.
  usemtl page_NN  -> texture index is looked up from the UV rectangle (bounding box of the
                     face UVs) on page NN; UVs must lie on rectangle corners
                     (the way geom_export writes them).
  usemtl tex_N    -> texture index N; UVs (if present) only choose the corner assignment,
                     without UVs the default assignment is used.
  usemtl flat_N   -> texture index N (this is how geom_export exports flat colours)
  otherwise       -> --default-tex
Quads are split as (0,1,2)+(2,3,0); with rectangle UVs this gives exactly the
original assignment (flags 0x800 / 0x810).
Vertex colours: "v x y z r g b" (0..1, 1.0 = original brightness 0x80); otherwise 0x80.
Collision: --coll-obj file.obj (polygons, groups "matNN...") or one convex collision polygon from each face (quad if planar), entered
into every grid cell the polygon (xz bounding box) touches. Material: group/object
"..._matNN" or --mat. (--no-coll: no collision)

Output: one model per occupied (cell, height band). All existing tiles are replaced
(except with --keep-tiles: old tiles + models stay, new ones are added).
Strats/paths/zones (rest of MAP) and all other chunks stay unchanged.
"""
import argparse
import math
import os
import re
import struct
import sys
from collections import defaultdict

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import geomlib as G  # noqa: E402

PATTERNS = [0x800, 0x810, 0x820, 0x830, 0x000, 0x020]


def read_obj(path, export_space=False):
    """Faces and points in export space (converted from file space unless export_space)."""
    V, C, VT = [], [], []
    faces = []   # (idx list, vt list or None, material, group)
    mat, grp = None, ''
    for line in open(path, encoding='latin1'):
        p = line.split()
        if not p:
            continue
        if p[0] == 'v':
            v = tuple(float(x) for x in p[1:4])
            V.append(v if export_space else G.to_file(v))
            C.append(tuple(float(x) for x in p[4:7]) if len(p) >= 7 else None)
        elif p[0] == 'vt':
            VT.append((float(p[1]), float(p[2])))
        elif p[0] == 'usemtl':
            mat = p[1] if len(p) > 1 else None
        elif p[0] in ('o', 'g'):
            grp = p[1] if len(p) > 1 else ''
        elif p[0] == 'f':
            vi, ti = [], []
            for tok in p[1:]:
                s = tok.split('/')
                a = int(s[0]); vi.append(a - 1 if a > 0 else len(V) + a)
                if len(s) > 1 and s[1]:
                    b = int(s[1]); ti.append(b - 1 if b > 0 else len(VT) + b)
            ti = ti if len(ti) == len(vi) else None
            if not export_space:
                vi = G.file_order(vi)
                ti = ti and G.file_order(ti)
            faces.append((vi, ti, mat, grp))
    return V, C, VT, faces


def find_tex(texs, page, uvs):
    """Texture index on page 'page' whose rectangle best matches the UVs."""
    us = [u * 256 for u, v in uvs]; vs = [v * 256 for u, v in uvs]
    best, bd = None, 1e9
    for i, t in enumerate(texs):
        if t['page'] != page or (t['flags'] & 1):
            continue
        d = (abs(t['u0'] - min(us)) + abs(t['u1'] + 1 - max(us)) + abs(t['v0'] - min(vs)) + abs(t['v1'] + 1 - max(vs)))
        if d < bd:
            best, bd = i, d
    return best, bd


def match_pattern(uvs, rect):
    """Best (flags, rotation) so that face_uvs(flags) hits the given corner UVs."""
    best = (1e9, 0x800, 0)
    for fl in PATTERNS:
        ref = G.face_uvs(fl, rect)
        for rot in range(3):
            r = [uvs[(k + rot) % 3] for k in range(3)]
            d = sum(abs(a[0] - b[0]) + abs(a[1] - b[1]) for a, b in zip(r, ref))
            if d < best[0]:
                best = (d, fl, rot)
    return best[1], best[2], best[0]


def color_u32(c):
    if c is None:
        return 0x01808080
    r, g, b = (max(0, min(255, int(round(x * 128)))) for x in c)
    return (1 << 24) | (b << 16) | (g << 8) | r


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('base'); ap.add_argument('obj'); ap.add_argument('out')
    ap.add_argument('--default-tex', type=int, default=None)
    ap.add_argument('--mat', type=int, default=0, help='collision material (default 0)')
    ap.add_argument('--no-coll', action='store_true')
    ap.add_argument('--coll-obj', help='collision from a separate OBJ (e.g. collision.obj from geom_export) '
                    'instead of the render faces')
    ap.add_argument('--keep-tiles', action='store_true')
    ap.add_argument('--zones', choices=['flat', 'keep'], default='flat',
                    help='flat: disable zone visibility (PVS), draw everything; keep: original')
    ap.add_argument('--band', type=float, default=6.0, help='height band per tile (units)')
    ap.add_argument('--export-space', action='store_true',
                    help='OBJ in raw export space (z < 0; files of geom_export --export-space / before 2026-11-07)')
    a = ap.parse_args()

    ch, wf, trak, mp = G.load_level(a.base)
    pages, texs = G.parse_text(G.wad_get(ch, 'TEXT'))
    W, H = mp['W'], mp['H']
    V, C, VT, faces = read_obj(a.obj, a.export_space)
    deftex = a.default_tex
    if deftex is None:
        deftex = next(i for i, t in enumerate(texs) if not (t['flags'] & 1) and t['page'] >= 0)

    # ---- triangles with texture + flags
    tris = []      # (pos[3], col[3], tex, flags)
    polys = []     # (pos list, mat) for collision
    warn = defaultdict(int)
    for vi, ti, mat, grp in faces:
        pts = [V[i] for i in vi]; cols = [C[i] for i in vi]
        uvs = [(VT[i][0], 1 - VT[i][1]) for i in ti] if ti else None
        m = re.search(r'mat(\d+)', grp or '')
        polys.append((pts, int(m.group(1)) if m else a.mat))
        tex = None
        mm = re.match(r'(page|tex|flat)_(\d+)$', mat or '')
        if mm and mm.group(1) == 'page' and uvs:
            tex, dist = find_tex(texs, int(mm.group(2)), uvs)
            if tex is None or dist > 2:
                warn['UV rectangle not found'] += 1
        elif mm:
            tex = int(mm.group(2))
        if tex is None:
            tex = deftex
        t = texs[tex]
        rect = G.tex_uv(t)
        for k in range(1, len(vi) - 1):
            idx = [0, k, k + 1]
            P = [pts[j] for j in idx]; Cc = [cols[j] for j in idx]
            if uvs and not (t['flags'] & 1):
                fl, rot, d = match_pattern([uvs[j] for j in idx], rect)
                if d > 2.5 / 256:
                    warn['UVs not on rectangle corners'] += 1
                P = [P[(j + rot) % 3] for j in range(3)]; Cc = [Cc[(j + rot) % 3] for j in range(3)]
            else:
                fl = 0x800 if k % 2 else 0x810
            tris.append((P, Cc, tex, fl))

    if a.coll_obj:
        cV, _, _, cfaces = read_obj(a.coll_obj, a.export_space)
        polys = []
        for vi, ti, mat, grp in cfaces:
            m = re.search(r'mat(\d+)', grp or '')
            polys.append(([cV[i] for i in vi], int(m.group(1)) if m else a.mat))

    # ---- cells/bands
    def key_of(x, y, z):
        cx, cz = math.floor(x), math.floor(-z)
        return (cx, cz, math.floor(y / a.band))

    cells = defaultdict(lambda: dict(tris=[], coll=[]))
    for tr in tris:
        cen = np.mean(tr[0], axis=0)
        cells[key_of(*cen)]['tris'].append(tr)
    if not a.no_coll:
        for pts, mat in polys:
            P = np.array(pts)
            cen = P.mean(0)
            band = math.floor(cen[1] / a.band)
            for cx in range(math.floor(P[:, 0].min()), math.floor(P[:, 0].max() - 1e-6) + 1):
                for cz in range(math.floor(-P[:, 2].max()), math.floor(-P[:, 2].min() - 1e-6) + 1):
                    cells[(cx, cz, band)]['coll'].append((pts, mat))

    # ---- build models
    models, tiles = [], []
    if a.keep_tiles:
        models = list(trak)
        for i in range(mp['N']):
            x, y, z, rot, mi, fl = mp['A'][i]
            tiles.append(dict(x=x, y=y, z=z, rot=rot, model=mi, cols=mp['cols'][i], flag=fl,
                              altcols=mp['altcols'].get(i, [])))
    for (cx, cz, band), cell in sorted(cells.items()):
        if not (0 <= cx < W and 0 <= cz < H):
            warn[f'outside the {W}x{H} grid'] += 1
            continue
        allp = [p for t in cell['tris'] for p in t[0]] or [p for c in cell['coll'] for p in c[0]]
        ty = float(round(np.median([p[1] for p in allp])))
        ox, oz = cx + 0.5, -(cz + 0.5)
        m = G.Model(); m.dirty = True
        vmap, cols = {}, []
        facelist = []
        for P, Cc, tex, fl in sorted(cell['tris'], key=lambda t: (texs[t[2]]['page'], texs[t[2]]['flags'], t[2])):
            ids = []
            for p, c in zip(P, Cc):
                lp = (round(p[0] - ox, 5), round(p[1] - ty, 5), round(p[2] - oz, 5))
                k = (lp, color_u32(c))
                if k not in vmap:
                    vmap[k] = len(m.verts); m.verts.append([*lp, 0.0, 0.0, 0.0]); cols.append(k[1])
                ids.append(vmap[k])
            if len(set(ids)) < 3:
                continue
            facelist.append((ids, tex, fl))
        prev = None
        nacc = np.zeros((len(m.verts), 3))
        for ids, tex, fl in facelist:
            plane = G.face_plane(*[m.verts[i][:3] for i in ids])
            st = (texs[tex]['page'], texs[tex]['flags'])
            # bit 0x8 = new render batch (reload state). Batch buffer 0x5d15f8 holds ~4800
            # output vertices -> restart at least every 256 faces.
            if st != prev or len(m.faces) % 256 == 0:
                fl |= 0x8
                prev = st
            m.faces.append(dict(flags=fl, v=tuple(ids), tex=tex, pad=0, plane=plane))
            for i in ids:
                nacc[i] += plane[:3]
        for i, v in enumerate(m.verts):
            n = nacc[i]; ln = np.linalg.norm(n)
            v[3:6] = (n / ln).tolist() if ln > 0 else [0.0, 1.0, 0.0]
        m.verts = [tuple(v) for v in m.verts]
        for pts, mat in cell['coll']:
            loc = [(p[0] - ox, p[1] - ty, p[2] - oz) for p in pts]
            if max(abs(c) for p in loc for c in p) > 7.9:
                warn['collision polygon > 7.9 units from tile centre (skipped)'] += 1
                continue
            recs = []
            if len(loc) == 4:
                r = G.coll_from_polygon(loc, mat)
                n = np.cross(np.subtract(loc[1], loc[0]), np.subtract(loc[2], loc[0]))
                if r is not None and abs(np.dot(n / (np.linalg.norm(n) or 1), np.subtract(loc[3], loc[0]))) < 1e-3:
                    recs = [r]
            if not recs:
                for k in range(1, len(loc) - 1):
                    r = G.coll_from_polygon([loc[0], loc[k], loc[k + 1]], mat)
                    if r is not None:
                        recs.append(r)
            for r in recs:
                ny = struct.unpack_from('<b', r, 3)[0]
                if ny < 0:
                    r = bytes([r[0] | 4]) + r[1:]
                m.coll[0].append(r)
        if len(m.verts) > 8000 or len(m.faces) > 65535:
            sys.exit('tile too large (transform buffer 0x58b0f0 holds 8000 vertices)')
        models.append(m)
        tiles.append(dict(x=ox, y=ty, z=oz, rot=0, model=len(models) - 1, cols=[cols]))

    if not models:
        sys.exit('no geometry')
    newtrak = G.build_trak(models)
    newmap = G.build_map_tiles(G.wad_get(ch, 'MAP'), wf, [len(m.verts) for m in trak], tiles, a.zones,
                                 mp['anim'] if a.keep_tiles else ())
    G.wad_set(ch, 'TRAK', newtrak)
    G.wad_set(ch, 'MAP', newmap)
    G.wad_write(a.out, ch)
    nf = sum(len(m.faces) for m in models); nc = sum(len(m.coll[0]) for m in models)
    print(f'{a.out}: {len(models)} models, {len(tiles)} tiles, {nf} triangles, {nc} collision polygons')
    for k, v in warn.items():
        print(f'  WARNING: {k}: {v}x')


if __name__ == '__main__':
    main()