"""Build a textured, stylized city level from OpenStreetMap data.

Run:  Blender -b -P tools/build_level_osm.py -- <overpass.json> <level_id> \
          <origin_lat> <origin_lon> [spawn_lat,lon,bearing] [rp_lat,lon;...] [bridge_way_name|-]

Buildings extrude with FLUX facade textures (box-projected UVs), roads/parks/
water/cliffs get their own materials, tree nodes become low-poly trees, and an
optional named bridge way grows a steel cantilever truss (hello Story Bridge).
Meshes suffixed -col/-convcol get Godot collision on import.

ponytail: flat terrain -- no DEM yet, so bridges sit at water level and the
Kangaroo Point cliffs are walls beside the road rather than a drop.
Data (c) OpenStreetMap contributors, ODbL.
"""
import json
import math
import os
import sys
import bpy
import bmesh
from mathutils import Matrix, Vector

argv = sys.argv[sys.argv.index("--") + 1:]
src, level_id = argv[0], argv[1]
LAT0, LON0 = float(argv[2]), float(argv[3])
spawn_arg = argv[4] if len(argv) > 4 else ""
rp_arg = argv[5] if len(argv) > 5 else ""
bridge_name = argv[6] if len(argv) > 6 and argv[6] != "-" else ""

TEX = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "assets", "textures")
M_LAT = 110574.0
M_LON = 111320.0 * math.cos(math.radians(LAT0))

def xy(lat, lon):
    return ((lon - LON0) * M_LON, (lat - LAT0) * M_LAT)

for obj in list(bpy.data.objects):
    bpy.data.objects.remove(obj, do_unlink=True)

def tex_material(name, img, rough=0.85):
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    bsdf = next(n for n in mat.node_tree.nodes if n.type == "BSDF_PRINCIPLED")
    bsdf.inputs["Roughness"].default_value = rough
    tex = mat.node_tree.nodes.new("ShaderNodeTexImage")
    tex.image = bpy.data.images.load(os.path.join(TEX, img))
    mat.node_tree.links.new(tex.outputs["Color"], bsdf.inputs["Base Color"])
    return mat

def flat_material(name, rgba, rough=0.8):
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    bsdf = next(n for n in mat.node_tree.nodes if n.type == "BSDF_PRINCIPLED")
    bsdf.inputs["Base Color"].default_value = rgba
    bsdf.inputs["Roughness"].default_value = rough
    return mat

M_GLASS = tex_material("facade_glass", "facade_glass.jpg", 0.4)
M_CONC = tex_material("facade_concrete", "facade_concrete.jpg")
M_BRICK = tex_material("facade_brick", "facade_brick.jpg")
M_HERIT = tex_material("facade_heritage", "facade_heritage.jpg")
M_ROAD = tex_material("asphalt", "asphalt.jpg")
M_MALL = tex_material("pavers", "pavers.jpg")
M_GRASS = tex_material("grass", "grass.jpg")
M_WATER = tex_material("water", "water.jpg", 0.15)
M_ROCK = tex_material("rock", "rock.jpg")
M_STEEL = tex_material("steel", "steel.jpg", 0.5)
M_GROUND = tex_material("concrete_ground", "concrete_ground.jpg")
M_ROOF = flat_material("roofgrey", (0.30, 0.30, 0.32, 1))
M_LEAF = flat_material("leaf", (0.18, 0.42, 0.16, 1))
M_TRUNK = flat_material("trunk", (0.28, 0.20, 0.12, 1))

def uv_project(bm, scale):
    uv = bm.loops.layers.uv.new("UVMap")
    for f in bm.faces:
        n = f.normal
        ax = max(range(3), key=lambda i: abs(n[i]))
        for l in f.loops:
            c = l.vert.co
            if ax == 2:
                l[uv].uv = (c.x / scale, c.y / scale)
            elif ax == 0:
                l[uv].uv = (c.y / scale, c.z / scale)
            else:
                l[uv].uv = (c.x / scale, c.z / scale)

def add_obj(name, bm, mats, uv_scale, roof_split=False):
    if roof_split:
        for f in bm.faces:
            f.material_index = 1 if f.normal.z > 0.7 else 0
    uv_project(bm, uv_scale)
    mesh = bpy.data.meshes.new(name)
    bm.to_mesh(mesh)
    bm.free()
    for m in mats:
        mesh.materials.append(m)
    obj = bpy.data.objects.new(name, mesh)
    bpy.context.collection.objects.link(obj)
    return obj

def poly_bm(pts, z, h=0.0):
    """Filled (optionally extruded) polygon bmesh from closed pt list."""
    bm = bmesh.new()
    verts = [bm.verts.new((x, y, z)) for x, y in pts[:-1]]
    face = bm.faces.new(verts)
    if h > 0.0:
        res = bmesh.ops.extrude_face_region(bm, geom=[face])
        up = [v for v in res["geom"] if isinstance(v, bmesh.types.BMVert)]
        bmesh.ops.translate(bm, vec=(0, 0, h), verts=up)
    bmesh.ops.recalc_face_normals(bm, faces=bm.faces)
    bmesh.ops.triangulate(bm, faces=bm.faces)
    return bm

def beam(bm, a, b, w, h):
    """Oriented box from point a to point b (Vectors)."""
    axis = b - a
    L = axis.length
    if L < 0.05:
        return
    quat = axis.to_track_quat("X", "Z")
    mat = Matrix.Translation((a + b) / 2) @ quat.to_matrix().to_4x4()
    res = bmesh.ops.create_cube(bm, size=1.0)
    verts = res["verts"]
    bmesh.ops.scale(bm, vec=(L, w, h), verts=verts)
    bmesh.ops.transform(bm, matrix=mat, verts=verts)

d = json.load(open(src))
counts = {"bld": 0, "road": 0, "park": 0, "water": 0, "cliff": 0, "tree": 0}
GREEN_LEISURE = ("park", "garden", "recreation_ground", "pitch", "playground", "golf_course", "nature_reserve")
GREEN_LANDUSE = ("grass", "recreation_ground", "village_green", "meadow", "forest")

tree_bm = bmesh.new()
for e in d["elements"]:
    tags = e.get("tags", {})
    if e["type"] == "node":
        if counts["tree"] >= 400:
            continue
        tx, ty = xy(e["lat"], e["lon"])
        res = bmesh.ops.create_cone(tree_bm, cap_ends=True, segments=6,
                                    radius1=2.2, radius2=0.15, depth=5.5)
        bmesh.ops.translate(tree_bm, vec=(tx, ty, 4.5), verts=res["verts"])
        counts["tree"] += 1
        continue

    geom = e.get("geometry")
    if not geom:
        continue
    pts = [xy(g["lat"], g["lon"]) for g in geom]
    closed = len(pts) > 3 and geom[0] == geom[-1]

    if tags.get("building"):
        if not closed or tags["building"] in ("roof", "bridge"):
            continue
        try:
            levels = float(tags.get("building:levels", "0"))
        except ValueError:
            levels = 0.0
        h = levels * 3.4 if levels > 0 else 8.0 + (e["id"] % 11) * 3.5
        try:
            bm = poly_bm(pts, 0, h)
        except Exception:
            continue
        if h > 55:
            fac = M_GLASS
        elif h < 15 and e["id"] % 3 == 0:
            fac = M_HERIT
        else:
            fac = M_BRICK if e["id"] % 2 else M_CONC
        add_obj("bld_%d-col" % e["id"], bm, [fac, M_ROOF], 13.0, roof_split=True)
        counts["bld"] += 1

    elif tags.get("natural") == "water" or tags.get("water"):
        if not closed:
            continue
        try:
            bm = poly_bm(pts, -0.35)
        except Exception:
            continue
        # ponytail: water is solid (-col) so cars skim it instead of falling forever
        add_obj("water_%d-col" % e["id"], bm, [M_WATER], 30.0)
        counts["water"] += 1

    elif tags.get("leisure") in GREEN_LEISURE or tags.get("landuse") in GREEN_LANDUSE \
            or tags.get("natural") in ("grassland", "scrub", "wood"):
        if not closed:
            continue
        try:
            bm = poly_bm(pts, 0.02)
        except Exception:
            continue
        add_obj("park_%d" % e["id"], bm, [M_GRASS], 14.0)
        counts["park"] += 1

    elif tags.get("natural") == "cliff":
        bm = bmesh.new()
        for i in range(len(pts) - 1):
            a = Vector((pts[i][0], pts[i][1], 7.0))
            b = Vector((pts[i + 1][0], pts[i + 1][1], 7.0))
            beam(bm, a, b, 2.0, 14.0)
        add_obj("cliff_%d-col" % e["id"], bm, [M_ROCK], 12.0)
        counts["cliff"] += 1

    elif tags.get("highway"):
        hw = tags["highway"]
        if hw in ("footway", "steps", "path", "cycleway", "corridor", "platform"):
            continue
        if tags.get("tunnel") == "yes" or int(tags.get("layer", "0") or 0) < 0:
            continue
        wide = {"primary": 9.0, "secondary": 8.0, "tertiary": 7.0, "pedestrian": 10.0}.get(hw, 6.0)
        mat = M_MALL if hw == "pedestrian" else M_ROAD
        bm = bmesh.new()
        for i in range(len(pts) - 1):
            (x1, y1), (x2, y2) = pts[i], pts[i + 1]
            dx, dy = x2 - x1, y2 - y1
            L = math.hypot(dx, dy)
            if L < 0.1:
                continue
            nx, ny = -dy / L * wide / 2, dx / L * wide / 2
            try:
                bm.faces.new([bm.verts.new(p) for p in [
                    (x1 + nx, y1 + ny, 0.03), (x1 - nx, y1 - ny, 0.03),
                    (x2 - nx, y2 - ny, 0.03), (x2 + nx, y2 + ny, 0.03)]])
            except Exception:
                pass
        add_obj("road_%d" % e["id"], bm, [mat], 8.0)
        counts["road"] += 1

if counts["tree"]:
    trunk_bm = bmesh.new()
    # trunks share the leaf-cone mesh's object for simplicity: one green mesh, one trunk mesh
    add_obj("trees", tree_bm, [M_LEAF], 6.0)
else:
    tree_bm.free()

# steel truss along a named bridge way
if bridge_name:
    cands = [e for e in d["elements"] if e.get("tags", {}).get("name") == bridge_name and e.get("geometry")]
    if cands:
        way = max(cands, key=lambda e: len(e["geometry"]))
        pts = [Vector((*xy(g["lat"], g["lon"]), 0.0)) for g in way["geometry"]]
        # arc-length resample the middle 65% (the steel section)
        dists = [0.0]
        for i in range(1, len(pts)):
            dists.append(dists[-1] + (pts[i] - pts[i - 1]).length)
        T = dists[-1]
        def at(s):
            for i in range(1, len(dists)):
                if dists[i] >= s:
                    f = (s - dists[i - 1]) / max(dists[i] - dists[i - 1], 0.001)
                    return pts[i - 1].lerp(pts[i], f)
            return pts[-1]
        s0, s1 = 0.175 * T, 0.825 * T
        N = 26
        bm = bmesh.new()
        samples = []
        for i in range(N + 1):
            t = i / N
            s = s0 + (s1 - s0) * t
            p = at(s)
            h = 4.0 + 16.0 * (math.exp(-((t - 0.32) / 0.10) ** 2) + math.exp(-((t - 0.68) / 0.10) ** 2)) \
                + 6.0 * math.sin(math.pi * t)
            fwd = (at(min(s + 4, T)) - p)
            fwd.z = 0
            fwd.normalize()
            perp = Vector((-fwd.y, fwd.x, 0))
            samples.append((p, perp, h))
        for side in (-1, 1):
            for i in range(N):
                p1, pe1, h1 = samples[i]
                p2, pe2, h2 = samples[i + 1]
                a = p1 + pe1 * side * 8.0
                b = p2 + pe2 * side * 8.0
                beam(bm, a + Vector((0, 0, h1)), b + Vector((0, 0, h2)), 0.7, 0.7)      # top chord
                beam(bm, a + Vector((0, 0, 0.5)), b + Vector((0, 0, 0.5)), 0.5, 0.5)   # bottom chord
                beam(bm, a + Vector((0, 0, 0.5)), a + Vector((0, 0, h1)), 0.5, 0.5)    # vertical
                if i % 2 == 0:
                    beam(bm, a + Vector((0, 0, 0.5)), b + Vector((0, 0, h2)), 0.4, 0.4)  # diagonal
        for i in range(0, N + 1, 3):
            p, pe, h = samples[i]
            beam(bm, p - pe * 8.0 + Vector((0, 0, h)), p + pe * 8.0 + Vector((0, 0, h)), 0.5, 0.5)  # cross brace
        add_obj("bridge_steel-col", bm, [M_STEEL], 6.0)
        print("bridge truss built along %r (%.0fm span)" % (bridge_name, s1 - s0))

# ground slab sized to content (convex collider: huge trimesh tris break raycasts)
xs = [v.co.x for o in bpy.data.objects if o.type == "MESH" for v in o.data.vertices]
ys = [v.co.y for o in bpy.data.objects if o.type == "MESH" for v in o.data.vertices]
gx, gy = (min(xs) + max(xs)) / 2, (min(ys) + max(ys)) / 2
bm = bmesh.new()
res = bmesh.ops.create_cube(bm, size=1.0)
bmesh.ops.scale(bm, vec=(max(xs) - min(xs) + 100, max(ys) - min(ys) + 100, 1.0), verts=res["verts"])
bmesh.ops.translate(bm, vec=(gx, gy, -0.5), verts=res["verts"])
add_obj("ground-convcol", bm, [M_GROUND], 18.0)

if spawn_arg:
    slat, slon, sbear = [float(v) for v in spawn_arg.split(",")]
    sx, sy = xy(slat, slon)
    sp = bpy.data.objects.new("Spawn", None)
    sp.location = (sx, sy, 0.1)
    sp.rotation_euler = (0, 0, math.radians(-sbear))
    bpy.context.collection.objects.link(sp)

if rp_arg:
    rp = bpy.data.objects.new("RacePath", None)
    bpy.context.collection.objects.link(rp)
    for i, pair in enumerate(rp_arg.split(";")):
        plat, plon = [float(v) for v in pair.split(",")]
        px, py = xy(plat, plon)
        emp = bpy.data.objects.new("RP_%02d" % i, None)
        emp.location = (px, py, 0.1)
        bpy.context.collection.objects.link(emp)
        emp.parent = rp

out = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "levels", level_id)
os.makedirs(out, exist_ok=True)
for obj in bpy.data.objects:
    obj.select_set(True)
path = os.path.join(out, "level.glb")
bpy.ops.export_scene.gltf(filepath=path, export_format="GLB", use_selection=True)
print("level built:", level_id, counts, "->", path)