"""Read-only shading audit for a Studio/bake job. Blender -- --job DIR --output DIR."""
import argparse
import hashlib
import json
import math
import sys
from pathlib import Path
import bpy
import numpy as np
from mathutils import Vector

p = argparse.ArgumentParser()
p.add_argument('--job', required=True)
p.add_argument('--output', required=True)
p.add_argument('--low-model')
p.add_argument('--quick', action='store_true')
p.add_argument('--normal-strength', type=float)
args = p.parse_args(sys.argv[sys.argv.index('--') + 1:])
job, out = Path(args.job).resolve(), Path(args.output).resolve()
out.mkdir(parents=True, exist_ok=True)


def load(path):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.gltf(filepath=str(path))
    meshes = [o for o in bpy.context.scene.objects if o.type == 'MESH']
    assert len(meshes) == 1
    obj = meshes[0]
    bpy.context.view_layer.objects.active = obj
    obj.select_set(True)
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    return obj


def describe(obj):
    result = []
    for mat in obj.data.materials:
        bsdf = next((n for n in mat.node_tree.nodes if n.type == 'BSDF_PRINCIPLED'), None)
        channels = {}
        if bsdf:
            for name in ['Base Color', 'Metallic', 'Roughness', 'Normal']:
                socket = bsdf.inputs[name]
                channels[name] = dict(linked=socket.is_linked,
                    source=socket.links[0].from_node.type if socket.is_linked else None)
        result.append(dict(name=mat.name, channels=channels,
            images=[dict(name=n.image.name, size=list(n.image.size), colorspace=n.image.colorspace_settings.name)
                    for n in mat.node_tree.nodes if n.type == 'TEX_IMAGE' and n.image]))
    return result


def setup(obj):
    pts = np.array([obj.matrix_world @ v.co for v in obj.data.vertices])
    lo, hi = pts.min(0), pts.max(0)
    center = Vector((lo + hi) / 2)
    h = hi[2] - lo[2]
    scene = bpy.context.scene
    scene.render.engine = 'CYCLES'
    scene.cycles.device = 'CPU'
    scene.cycles.samples = 16
    scene.cycles.use_denoising = True
    scene.render.resolution_x = 800
    scene.render.resolution_y = 800
    scene.render.resolution_percentage = 100
    scene.render.image_settings.file_format = 'PNG'
    scene.world = bpy.data.worlds.new('AuditWorld')
    scene.world.use_nodes = True
    scene.world.node_tree.nodes['Background'].inputs['Color'].default_value = (.4, .4, .4, 1)
    scene.world.node_tree.nodes['Background'].inputs['Strength'].default_value = .6
    for name, xyz, power, size in [('Key', (-2, 3, 3), 180, 2), ('Fill', (2, -2, 1.5), 100, 2)]:
        data = bpy.data.lights.new(name, 'AREA'); data.energy = power; data.shape = 'DISK'; data.size = size*h
        light = bpy.data.objects.new(name, data); scene.collection.objects.link(light)
        light.location = center + Vector(xyz)*h
        light.rotation_euler = (center-light.location).to_track_quat('-Z', 'Y').to_euler()
    camera = bpy.data.objects.new('Camera', bpy.data.cameras.new('Camera'))
    scene.collection.objects.link(camera); scene.camera = camera
    camera.data.type = 'ORTHO'; camera.data.ortho_scale = max(h*1.17, (hi[0]-lo[0])*1.17)
    camera.location = center + Vector((1.7, 3, 1))*h
    camera.rotation_euler = (center-camera.location).to_track_quat('-Z', 'Y').to_euler()
    scene.view_settings.view_transform = 'Standard'


def capture(name):
    bpy.context.scene.render.filepath = str(out / (name+'.png'))
    bpy.ops.render.render(write_still=True)


def normal_quality(obj):
    node = next((n for mat in obj.data.materials for n in mat.node_tree.nodes
                 if n.type == 'TEX_IMAGE' and n.image and n.image.name.startswith('normal')), None)
    if node is None:
        return {}
    image = node.image
    w, h = image.size[:]
    pixels = np.empty(w*h*4, dtype=np.float32)
    image.pixels.foreach_get(pixels)
    pixels = pixels.reshape(h, w, 4)
    coverage = np.zeros((h, w), dtype=bool)
    uv_data = obj.data.uv_layers.active.data
    obj.data.calc_loop_triangles()
    corners = []
    for tri in obj.data.loop_triangles:
        uvs = np.array([uv_data[i].uv[:] for i in tri.loops])
        for vi, uv in zip(tri.vertices, uvs):
            corners.append(tuple(np.round(list(obj.data.vertices[vi].co)+list(uv), 6)))
        x, y = uvs[:, 0]*w-.5, uvs[:, 1]*h-.5
        x0, x1 = max(0, int(np.floor(x.min()))), min(w-1, int(np.ceil(x.max())))
        y0, y1 = max(0, int(np.floor(y.min()))), min(h-1, int(np.ceil(y.max())))
        den = (y[1]-y[2])*(x[0]-x[2])+(x[2]-x[1])*(y[0]-y[2])
        if abs(den)<1e-8 or x1<x0 or y1<y0:
            continue
        xx, yy = np.meshgrid(np.arange(x0,x1+1),np.arange(y0,y1+1))
        a=((y[1]-y[2])*(xx-x[2])+(x[2]-x[1])*(yy-y[2]))/den
        b=((y[2]-y[0])*(xx-x[2])+(x[0]-x[2])*(yy-y[2]))/den
        coverage[y0:y1+1,x0:x1+1] |= (a>=0)&(b>=0)&(a+b<=1)
    length = np.linalg.norm(pixels[:,:,:3]*2-1, axis=2)
    invalid = coverage & ((length < .85) | (length > 1.15))
    return dict(uvCoveredPixels=int(coverage.sum()),
        invalidNormalDirectionPixels=int(invalid.sum()),
        invalidNormalDirectionPercent=float(invalid.sum()/max(1,coverage.sum())*100),
        nearlyBlackNormalPixels=int((coverage & (pixels[:,:,:3].max(axis=2)<.03)).sum()),
        oppositeHemispherePixels=int((coverage & (pixels[:,:,2]<.5)).sum()),
        geometryUvFingerprint=hashlib.sha256(np.array(sorted(corners),dtype='<f4').tobytes()).hexdigest())


report = {}
items = [('low', Path(args.low_model).resolve() if args.low_model else job/'model_baked.glb')]
if not args.quick:
    items.insert(0,('high',job/'model.glb'))
for tag, filename in items:
    obj = load(filename)
    if args.normal_strength is not None and tag=='low':
        for mat in obj.data.materials:
            for node in mat.node_tree.nodes:
                if node.type=='NORMAL_MAP':
                    node.inputs['Strength'].default_value=args.normal_strength
    report[tag] = dict(triangles=sum(len(poly.vertices)-2 for poly in obj.data.polygons), materials=describe(obj),
                       normalQuality=normal_quality(obj) if tag=='low' else None)
    setup(obj)
    capture(tag+'_lit_full')
    if args.quick:
        if args.normal_strength is not None:
            bpy.ops.object.select_all(action='DESELECT');obj.select_set(True)
            bpy.context.view_layer.objects.active=obj
            bpy.ops.export_scene.gltf(filepath=str(out/'normal_strength_preview.glb'),export_format='GLB',use_selection=True,export_animations=False)
        continue
    if tag == 'low':
        for mat in obj.data.materials:
            nt = mat.node_tree; bsdf = next(n for n in nt.nodes if n.type == 'BSDF_PRINCIPLED')
            for link in list(bsdf.inputs['Normal'].links): nt.links.remove(link)
        capture('low_without_normal')
        for mat in obj.data.materials:
            nt=mat.node_tree; bsdf=next(n for n in nt.nodes if n.type=='BSDF_PRINCIPLED')
            output=next(n for n in nt.nodes if n.type=='OUTPUT_MATERIAL')
            emit=nt.nodes.new('ShaderNodeEmission')
            base=bsdf.inputs['Base Color']
            if base.is_linked:nt.links.new(base.links[0].from_socket,emit.inputs['Color'])
            else:emit.inputs['Color'].default_value=base.default_value
            for link in list(output.inputs['Surface'].links):nt.links.remove(link)
            nt.links.new(emit.outputs['Emission'],output.inputs['Surface'])
        capture('low_albedo_unlit')
        # Plain material isolates the exported geometric/corner normals.
        mat=bpy.data.materials.new('Clay');mat.use_nodes=True
        bsdf=mat.node_tree.nodes.get('Principled BSDF');bsdf.inputs['Base Color'].default_value=(.5,.5,.5,1);bsdf.inputs['Roughness'].default_value=.8
        obj.data.materials.clear();obj.data.materials.append(mat)
        capture('low_clay_exported_normals')
        try:bpy.ops.mesh.customdata_custom_splitnormals_clear()
        except Exception:pass
        for poly in obj.data.polygons:poly.use_smooth=True
        capture('low_clay_recomputed_normals')
(out/'audit.json').write_text(json.dumps(report,indent=2),encoding='utf-8')
print('AUDIT_DONE',out)