"""Export the existing 37 accepted district placements as four spatial REPLACE LOD cells.

Run with Blender: blender --background --factory-startup --python-exit-code 1 \
    --python tools/city-massing/assemble-lod-district.py -- \
    --district tools/city-massing/district --grid 2
The original single-tile output remains independent and unchanged. Use a larger
`--grid` for broad procedural districts so Cesium can cull distant cells.
"""
import json
import math
import struct
import argparse
import sys
from pathlib import Path

import bpy
from mathutils import Vector

ROOT = Path(__file__).resolve().parents[2]
parser = argparse.ArgumentParser()
parser.add_argument('--district', default=str(ROOT / 'tools/city-massing/district'))
parser.add_argument('--output', default=None)
parser.add_argument('--grid', type=int, default=2)
parser.add_argument('--coarse-grid', type=int, default=0)
args = parser.parse_args(sys.argv[sys.argv.index('--') + 1:] if '--' in sys.argv else [])
if args.grid < 2:
    raise RuntimeError('--grid must be at least 2')
if args.coarse_grid and (args.coarse_grid < 2 or args.coarse_grid >= args.grid):
    raise RuntimeError('--coarse-grid must be at least 2 and smaller than --grid')
DISTRICT = Path(args.district).resolve()
OUT = Path(args.output).resolve() if args.output else DISTRICT / 'lod-preview'
LIBRARY_DIR = ROOT / 'tools/city-massing/house-library'
plan = json.loads((DISTRICT / 'seated-plan.json').read_text(encoding='utf-8'))
lib = json.loads((LIBRARY_DIR / 'library.json').read_text(encoding='utf-8'))
entries = lib['archetypes']
# Same asset-fit contract the placement generator used, so the two stages cannot disagree about
# what counts as an acceptable match.
ASSET_FIT = json.loads((ROOT / 'tools/city-massing/asset-fit.json').read_text(encoding='utf-8'))['variants']


def select_asset(building):
    family = building.get('archetype') or 'house'
    prefix = {'house': 'house', 'apartment': 'slab', 'midrise': 'midrise'}.get(family, 'house')
    # A building may name the exact approved variants it was planned for; otherwise any variant of
    # its family that fits the shared scale window is eligible.
    preferred = building.get('preferredAssets') or []
    family_key = family if family in ASSET_FIT else prefix
    limits = ASSET_FIT.get(family_key) or ASSET_FIT.get(prefix) or ASSET_FIT['house']
    min_x, min_y = limits['minScale']
    max_x, max_y = limits['maxScale']
    candidates = []
    for entry in entries:
        if preferred:
            if entry['id'] not in preferred:
                continue
        elif not entry['id'].startswith(prefix):
            continue
        width, depth = entry['nominal_footprint_m']
        sx, sy = building['width'] / width, building['depth'] / depth
        if not (min_x <= sx <= max_x and min_y <= sy <= max_y):
            continue
        preference = preferred.index(entry['id']) if entry['id'] in preferred else len(preferred) + 1
        score = preference * 0.25 + abs(entry['floors'] - building['floors']) * 0.8 + abs(math.log(sx)) + abs(math.log(sy))
        candidates.append((score, entry, sx, sy))
    return min(candidates, key=lambda item: item[0]) if candidates else None


placements, rejected = [], []
for building in plan['buildings']:
    choice = select_asset(building)
    if choice:
        _, entry, sx, sy = choice
        placements.append((building, entry, sx, sy))
    else:
        rejected.append(building['id'])
if not placements or len({b['id'] for b, _, _, _ in placements}) != len(placements):
    raise RuntimeError('The seated plan does not yield unique, nonempty placements')
source_plan = json.loads((DISTRICT / 'plan.json').read_text(encoding='utf-8'))
# Plan v4 lists modules under `buildings` and their roots under `sources`; older plans list only
# source footprints. Accept either so a module always resolves to a recorded source.
sources = {source['id'] for source in source_plan.get('buildings', source_plan.get('roofs', []))}
sources.update(source['id'] for source in source_plan.get('sources', []))
sources.update(building.get('sourceId') for building in source_plan.get('buildings', []) if building.get('sourceId'))
if any(b['sourceId'] not in sources or b['foundationBottom'] >= b['positionENU'][2]
       for b, _, _, _ in placements):
    raise RuntimeError('A placed building has no original source or invalid foundation seating')

# Partition placements into deterministic spatial cells. The original 2x2 names are kept
# for existing district and estate artifacts; larger grids use indexed cell names.
if args.grid == 2:
    xs = sorted(b['positionENU'][0] for b, _, _, _ in placements)
    ys = sorted(b['positionENU'][1] for b, _, _, _ in placements)
    mid_x, mid_y = xs[len(xs)//2], ys[len(ys)//2]
    cells = {name: [] for name in ('south-west', 'south-east', 'north-west', 'north-east')}
    for placement in placements:
        x, y, _ = placement[0]['positionENU']
        cell_id = ('north' if y >= mid_y else 'south') + ('-east' if x >= mid_x else '-west')
        cells[cell_id].append(placement)
else:
    xs = [b['positionENU'][0] for b, _, _, _ in placements]
    ys = [b['positionENU'][1] for b, _, _, _ in placements]
    min_x, max_x = min(xs), max(xs)
    min_y, max_y = min(ys), max(ys)
    span_x = max(max_x - min_x, 1e-6)
    span_y = max(max_y - min_y, 1e-6)
    cells = {}
    for placement in placements:
        x, y, _ = placement[0]['positionENU']
        ix = min(args.grid - 1, int((x - min_x) / span_x * args.grid))
        iy = min(args.grid - 1, int((y - min_y) / span_y * args.grid))
        cell_id = f'cell-{iy:02d}-{ix:02d}'
        cells.setdefault(cell_id, []).append(placement)

cells = {name: group for name, group in cells.items() if group}
if len(cells) < 2:
    raise RuntimeError('Not enough geographically separate building groups to demonstrate spatial LOD')

bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.context.scene.unit_settings.system = 'METRIC'
bpy.context.scene.unit_settings.scale_length = 1
OUT.mkdir(parents=True, exist_ok=True)
for stale in OUT.glob('*-lod0.glb'):
    stale.unlink()
for stale in OUT.glob('*-lod1.glb'):
    stale.unlink()
for stale in (OUT / 'root-anchor.glb', OUT / 'tileset.json', OUT / 'manifest.json'):
    if stale.exists():
        stale.unlink()
stone = bpy.data.materials.new('Buried_foundation_stone')
stone.diffuse_color = (0.32, 0.33, 0.32, 1)
stone.use_nodes = True
stone.node_tree.nodes['Principled BSDF'].inputs['Base Color'].default_value = stone.diffuse_color
stone.node_tree.nodes['Principled BSDF'].inputs['Roughness'].default_value = 0.92

# The dimensions/height vary per foundation; scale a shared unit cube mesh per instance.
foundation_mesh = bpy.data.meshes.new('Seated_foundation_unit_cube')
foundation_mesh.from_pydata(
    [(-.5, -.5, -.5), (.5, -.5, -.5), (.5, .5, -.5), (-.5, .5, -.5),
     (-.5, -.5, .5), (.5, -.5, .5), (.5, .5, .5), (-.5, .5, .5)],
    [],
    [(3, 2, 1, 0), (4, 5, 6, 7), (0, 1, 5, 4),
     (1, 2, 6, 5), (2, 3, 7, 6), (3, 0, 4, 7)],
)
foundation_mesh.materials.append(stone)
# Load the eight placed library templates together so common materials retain their names.
required = {entry['joined_object'] for _, entry, _, _ in placements}
required.update(entry['lod1']['joined_object'] for _, entry, _, _ in placements)
with bpy.data.libraries.load(str(LIBRARY_DIR / 'house-library.blend'), link=False) as (src, dst):
    missing = required.difference(src.objects)
    if missing:
        raise RuntimeError('Missing Blender library templates: ' + ', '.join(sorted(missing)))
    dst.objects = sorted(required)
templates = {obj.name: obj for obj in dst.objects if obj is not None}
if len(templates) != len(required) or any(obj.type != 'MESH' or not obj.data.polygons
                                          for obj in templates.values()):
    raise RuntimeError('A required Blender template is empty')
red_roof_sources = {source['id'] for source in source_plan.get('roofs', [])}
red_roof_sources.update({
    'osm-1285336668', 'osm-1285336669', 'osm-1285336671',
    'osm-1285336672', 'osm-1285336673', 'osm-1285336674',
    'osm-1285336675', 'trace-east-01', 'trace-east-02',
    'trace-east-04', 'trace-east-07', 'trace-east-09', 'trace-east-10',
})

# Scene variety without random re-rolls: rotate the shared library palettes per source
# building, so neighbouring buildings differ while every part of one building stays coherent.
ROOF_TONES = ('Roof_Slate', 'Roof_Slate_Light', 'Roof_Terracotta', 'Roof_Terra_Light')
PLASTER_TONES = ('Plaster_Cream', 'Plaster_Ivory', 'Plaster_Sand')


def source_seed(source_id):
    value = 7
    for character in str(source_id):
        value = (value * 131 + ord(character)) % 1000003
    return value


def material_remap(source_id):
    seed = source_seed(source_id)
    roof_shift = (seed // 3) % len(ROOF_TONES)
    wall_shift = seed % len(PLASTER_TONES)
    remap = {name: ROOF_TONES[(index + roof_shift) % len(ROOF_TONES)]
             for index, name in enumerate(ROOF_TONES)}
    remap.update({name: PLASTER_TONES[(index + wall_shift) % len(PLASTER_TONES)]
                  for index, name in enumerate(PLASTER_TONES)})
    remap['Roof_Membrane'] = ROOF_TONES[roof_shift]
    return remap


def template(entry, level):
    name = entry['joined_object'] if level == 'lod0' else entry['lod1']['joined_object']
    return templates[name]


def make_cell(group, level):
    collection = bpy.data.collections.new('District_' + level)
    bpy.context.scene.collection.children.link(collection)
    for building, entry, sx, sy in group:
        obj = template(entry, level).copy()  # Object copy shares mesh data with sibling placements.
        obj.name = building['id']
        collection.objects.link(obj)
        obj.location = building['positionENU']
        obj.rotation_euler = (0, 0, building['heading'])
        obj.scale = (sx, sy, 1)
        obj['source_footprint'] = building['sourceId']
        remap = material_remap(building['sourceId'])
        if building['sourceId'] in red_roof_sources:
            # This river-side row has red/brown roofs in the checked satellite mosaic.
            remap.update({'Roof_Slate': 'Roof_Terracotta', 'Roof_Slate_Light': 'Roof_Terra_Light'})
        for slot in obj.material_slots:
            if slot.material and slot.material.name in remap:
                material = bpy.data.materials.get(remap[slot.material.name])
                if material:
                    slot.link = 'OBJECT'
                    slot.material = material
        bottom = building['foundationBottom']
        x, y, top = building['positionENU']
        foundation = bpy.data.objects.new(building['id'] + '_buried_foundation', foundation_mesh)
        collection.objects.link(foundation)
        foundation.location = (x, y, (bottom + top) / 2)
        foundation.rotation_euler.z = building['heading']
        foundation.scale = (building['width'], building['depth'], top - bottom)
    bpy.context.view_layer.update()
    objects = list(collection.objects)
    corners = [obj.matrix_world @ Vector(corner) for obj in objects for corner in obj.bound_box]
    bounds = [[min(p[axis] for p in corners) for axis in range(3)],
              [max(p[axis] for p in corners) for axis in range(3)]]
    return collection, objects, bounds


def box(bounds):
    low, high = bounds
    centre = [(a + b) / 2 for a, b in zip(low, high)]
    half = [(b - a) / 2 for a, b in zip(low, high)]
    return [centre[0], centre[1], centre[2], half[0], 0, 0, 0, half[1], 0, 0, 0, half[2]]


def export_cell(path, objects, group, level):
    bpy.ops.object.select_all(action='DESELECT')
    for obj in objects:
        obj.select_set(True)
    bpy.context.view_layer.objects.active = objects[0]
    bpy.ops.export_scene.gltf(filepath=str(path), export_format='GLB', use_selection=True,
                              export_yup=True, export_materials='EXPORT', export_extras=True)
    with path.open('rb') as stream:
        header = stream.read(20)
        if len(header) != 20 or header[:4] != b'glTF':
            raise RuntimeError('Invalid GLB export ' + str(path))
        _, version, length, chunk_length, chunk_type = struct.unpack('<4s4I', header)
        if version != 2 or length != path.stat().st_size or chunk_type != 0x4E4F534A:
            raise RuntimeError('Invalid GLB header ' + str(path))
        doc = json.loads(stream.read(chunk_length))
    expected = {b['id'] for b, _, _, _ in group}
    names = [n.get('name') for n in doc['nodes'] if 'mesh' in n]
    if len(names) != 2 * len(expected) or sorted(names) != sorted(
            list(expected) + [name + '_buried_foundation' for name in expected]):
        raise RuntimeError('GLB mesh nodes do not match placed buildings and foundations: ' + str(path))
    if set(doc['scenes'][doc.get('scene', 0)]['nodes']) != set(range(len(doc['nodes']))):
        raise RuntimeError('Unexpected parent transform in ' + str(path))
    nodes = {node['name']: node for node in doc['nodes']}
    for building, _, sx, sy in group:
        x, y, top = building['positionENU']
        bottom = building['foundationBottom']
        for name, centre, scale in (
            (building['id'], top, (sx, 1, sy)),
            (building['id'] + '_buried_foundation', (bottom + top) / 2,
             (building['width'], top - bottom, building['depth'])),
        ):
            node = nodes[name]
            expected_pos = (x, centre, -y)  # Blender Z-up to glTF Y-up, only once.
            if ('matrix' in node or any(abs(a - b) > .002 for a, b in zip(
                    node['translation'], expected_pos)) or
                    any(abs(a - b) > .002 for a, b in zip(node['scale'], scale))):
                raise RuntimeError('Incorrect seated transform for ' + name + ' in ' + str(path))
    def triangles(mesh):
        return sum(doc['accessors'][part['indices']]['count'] // 3 if 'indices' in part
                   else doc['accessors'][part['attributes']['POSITION']]['count'] // 3
                   for part in mesh['primitives'])
    count = sum(triangles(doc['meshes'][node['mesh']]) for node in doc['nodes'] if 'mesh' in node)
    reference = sum((entry['validation'] if level == 'lod0' else entry['lod1']['validation'])['triangles'] + 12
                    for _, entry, _, _ in group)
    if count != reference:
        raise RuntimeError(f'{path.name}: {count} triangles, expected {reference}')
    meshes = {node['mesh'] for node in doc['nodes'] if 'mesh' in node}
    if len(expected) > 1 and len(meshes) >= len(names):
        raise RuntimeError('No shared mesh data in ' + str(path))
    return {'uri': path.name, 'bytes': path.stat().st_size, 'triangles': count,
            'meshNodes': len(names), 'uniqueMeshes': len(meshes)}


hierarchical = bool(args.coarse_grid)
coarse_groups, fine_coarse = {}, {}
if hierarchical:
    for cell_id, group in cells.items():
        x, y, _ = group[0][0]['positionENU']
        ix = min(args.coarse_grid - 1, int((x - min_x) / span_x * args.coarse_grid))
        iy = min(args.coarse_grid - 1, int((y - min_y) / span_y * args.coarse_grid))
        coarse_id = f'coarse-{iy:02d}-{ix:02d}'
        fine_coarse[cell_id] = coarse_id
        coarse_groups.setdefault(coarse_id, []).extend(group)

fine_tiles = {}
manifest_cells = []
all_bounds = []
for cell_id, group in cells.items():
    exports, bounds_per_level = {}, []
    for level in ('lod1', 'lod0'):
        collection, objects, bounds = make_cell(group, level)
        path = OUT / f'{cell_id}-{level}.glb'
        exports[level] = export_cell(path, objects, group, level)
        bounds_per_level.append(bounds)
        bpy.data.collections.remove(collection, do_unlink=True)
        for obj in objects:
            bpy.data.objects.remove(obj, do_unlink=True)
    bounds = [[min(b[0][axis] for b in bounds_per_level) for axis in range(3)],
              [max(b[1][axis] for b in bounds_per_level) for axis in range(3)]]
    all_bounds.append(bounds)
    volume = {'box': box(bounds)}
    if hierarchical:
        fine_tiles[cell_id] = {'boundingVolume': volume, 'geometricError': 0,
                               'content': {'uri': exports['lod0']['uri']}}
    else:
        fine_tiles[cell_id] = {'boundingVolume': volume, 'geometricError': 8,
                               'refine': 'REPLACE', 'content': {'uri': exports['lod1']['uri']},
                               'children': [{'boundingVolume': volume, 'geometricError': 0,
                                             'content': {'uri': exports['lod0']['uri']}}]}
    manifest_cells.append({'id': cell_id, 'count': len(group),
                           'sourceIds': sorted({b['sourceId'] for b, _, _, _ in group}),
                           'buildingIds': [b['id'] for b, _, _, _ in group],
                           'boundsENU': {'min': bounds[0], 'max': bounds[1]}, **exports})

if hierarchical:
    cell_tiles = []
    manifest_coarse = []
    for coarse_id, group in coarse_groups.items():
        collection, objects, bounds = make_cell(group, 'lod1')
        path = OUT / f'{coarse_id}-lod1.glb'
        coarse_export = export_cell(path, objects, group, 'lod1')
        bpy.data.collections.remove(collection, do_unlink=True)
        for obj in objects:
            bpy.data.objects.remove(obj, do_unlink=True)
        children = [fine_tiles[cell_id] for cell_id, parent in fine_coarse.items() if parent == coarse_id]
        cell_tiles.append({'boundingVolume': {'box': box(bounds)}, 'geometricError': 16,
                           'refine': 'REPLACE', 'content': {'uri': coarse_export['uri']},
                           'children': children})
        manifest_coarse.append({'id': coarse_id, 'count': len(group), **coarse_export})
else:
    cell_tiles = list(fine_tiles.values())
    manifest_coarse = []

root_bounds = [[min(b[0][axis] for b in all_bounds) for axis in range(3)],
               [max(b[1][axis] for b in all_bounds) for axis in range(3)]]
# The explicit root must be renderable to keep Cesium traversing its spatial children.
anchor_json = json.dumps({'asset': {'version': '2.0'}, 'scene': 0, 'scenes': [{'nodes': []}]},
                         separators=(',', ':')).encode('utf-8')
anchor_json += b' ' * ((-len(anchor_json)) % 4)
anchor = b'glTF' + struct.pack('<II', 2, 20 + len(anchor_json))
anchor += struct.pack('<I', len(anchor_json)) + b'JSON' + anchor_json
(OUT / 'root-anchor.glb').write_bytes(anchor)
tileset = {'asset': {'version': '1.1', 'generator': 'Blender scoped district LOD'},
           'geometricError': 800,
           'root': {'transform': plan['transform'], 'boundingVolume': {'box': box(root_bounds)},
                    'geometricError': 200, 'refine': 'REPLACE',
                     'content': {'uri': 'root-anchor.glb'}, 'children': cell_tiles}}
lod1_records = manifest_coarse if hierarchical else [cell['lod1'] for cell in manifest_cells]
lod0_records = [cell['lod0'] for cell in manifest_cells]
manifest = {'scope': plan.get('scope', 'evidence-bound district sample; not full-city reconstruction'),
            'origin': plan['origin'], 'notice': plan['notice'], 'count': len(placements),
            'cellCount': len(cells), 'cellIds': list(cells),
             'hierarchy': {'enabled': hierarchical, 'coarseGrid': args.coarse_grid if hierarchical else None,
                           'coarseCellCount': len(manifest_coarse)},
            'excludedUnfittedIds': rejected,
            'boundsENU': {'min': root_bounds[0], 'max': root_bounds[1]},
            'totals': {level: {'bytes': sum(item['bytes'] for item in (lod1_records if level == 'lod1' else lod0_records)),
                                'triangles': sum(item['triangles'] for item in (lod1_records if level == 'lod1' else lod0_records))}
                         for level in ('lod1', 'lod0')},
                                   'coarseCells': manifest_coarse,
             'cells': manifest_cells}
(OUT / 'tileset.json').write_text(json.dumps(tileset, indent=2), encoding='utf-8')
(OUT / 'manifest.json').write_text(json.dumps(manifest, indent=2), encoding='utf-8')
print('LOD_DISTRICT', json.dumps({'count': manifest['count'], 'cells': manifest['cellCount'],
                                   'totals': manifest['totals'], 'output': str(OUT)}))