#!/usr/bin/env python3
"""OBJ -> neue Weltgeometrie (TRAK + Kachelteil von MAP) in einer Kopie eines Level-WADs.

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

Eingabe-OBJ im Exportraum von geom_export.py (y oben, 1 Einheit = 1 Kachel, z < 0 liegt im
Raster; Rasterzelle (cx,cz) = (floor(x), floor(-z)), muss in 0..W-1 / 0..H-1 liegen).

Texturen: Die Engine kennt KEINE freien UVs. Jedes Dreieck zeigt eine Haelfte eines
Texturrechtecks (Eintrag der TEXT-Tabelle), die Ecken werden ueber Face-Flags gewaehlt.
  usemtl page_NN  -> Texturindex wird aus dem UV-Rechteck (Bounding-Box der Face-UVs)
                     auf Seite NN gesucht; UVs muessen auf Rechteck-Ecken liegen
                     (so wie geom_export sie schreibt).
  usemtl tex_N    -> Texturindex N; UVs (falls vorhanden) waehlen nur die Ecken-Belegung,
                     ohne UVs wird die Standardbelegung genommen.
  usemtl flat_N   -> Texturindex N (so exportiert geom_export Flachfarben)
  sonst           -> --default-tex
Vierecke werden als (0,1,2)+(2,3,0) zerlegt; mit Rechteck-UVs ergibt das genau die
Original-Belegung (Flags 0x800 / 0x810).
Vertexfarben: "v x y z r g b" (0..1, 1.0 = Originalhelligkeit 0x80); sonst 0x80.
Kollision: --coll-obj datei.obj (Polygone, Gruppen "matNN...") oder aus jedem Face ein konvexes Kollisionspolygon (Viereck wenn planar), eingetragen
in jede Rasterzelle, die das Polygon (xz-Bounding-Box) beruehrt. Material: Gruppe/Objekt
"..._matNN" bzw. --mat. (--no-coll: keine Kollision)

Ausgabe: Ein Modell pro belegter (Zelle, Hoehenband). Alle vorhandenen Kacheln werden ersetzt
(ausser --keep-tiles: alte Kacheln + Modelle bleiben, neue kommen dazu).
Strats/Pfade/Zonen (Rest von MAP) und alle anderen Chunks bleiben unveraendert.
"""
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):
    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.append(tuple(float(x) for x in p[1:4]))
            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)
            faces.append((vi, ti if len(ti) == len(vi) else None, mat, grp))
    return V, C, VT, faces


def find_tex(texs, page, uvs):
    """Texturindex auf Seite 'page', dessen Rechteck am besten zu den UVs passt."""
    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):
    """Beste (Flags, Rotation) so dass face_uvs(flags) die gegebenen Ecken-UVs trifft."""
    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='Kollisionsmaterial (Standard 0)')
    ap.add_argument('--no-coll', action='store_true')
    ap.add_argument('--coll-obj', help='Kollision aus eigenem OBJ (z.B. collision.obj von geom_export) '
                    'statt aus den Render-Faces')
    ap.add_argument('--keep-tiles', action='store_true')
    ap.add_argument('--zones', choices=['flat', 'keep'], default='flat',
                    help='flat: Zonen-Sichtbarkeit (PVS) abschalten, alles zeichnen; keep: Original')
    ap.add_argument('--band', type=float, default=6.0, help='Hoehenband pro Kachel (Einheiten)')
    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)
    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)

    # ---- Dreiecke mit Textur + Flags
    tris = []      # (pos[3], col[3], tex, flags)
    polys = []     # (pos list, mat) fuer Kollision
    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-Rechteck nicht gefunden'] += 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 nicht auf Rechteck-Ecken'] += 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)
        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))

    # ---- Zellen/Baender
    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))

    # ---- Modelle bauen
    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'ausserhalb des Rasters {W}x{H}'] += 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 = neuer Render-Batch (Zustand neu laden). Batch-Puffer 0x5d15f8 fasst ~4800
            # Ausgabe-Vertices -> spaetestens alle 256 Faces neu starten.
            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['Kollisionspolygon > 7.9 Einheiten vom Kachelzentrum (uebersprungen)'] += 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('Kachel zu gross (Transformpuffer 0x58b0f0 fasst 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('keine Geometrie')
    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)} Modelle, {len(tiles)} Kacheln, {nf} Dreiecke, {nc} Kollisionspolygone')
    for k, v in warn.items():
        print(f'  WARNUNG: {k}: {v}x')


if __name__ == '__main__':
    main()