"""Original UV image-painting study for Blender 5 Eevee
Run: blender --background --python render-texture-paint-study.py -- /new/output
Authored numeric image paint, no brush operator recording or AI inference
"""

import hashlib
import json
import math
import struct
import sys
import zlib
from pathlib import Path

import bpy

out = Path(sys.argv[sys.argv.index("--") + 1])
out.mkdir(parents=True, exist_ok=False)
w, h = 512, 256
recipe = {
    "kind": "Original numeric UV image paint, not captured viewport brush strokes",
    "width": w,
    "height": h,
    "baseRGBA": [38, 104, 114, 255],
    "band": {"centerV": 0.5, "halfWidth": 0.065, "feather": 0.008, "RGBA": [230, 127, 49, 255]},
    "dots": [{"u": 0.25, "v": 0.70, "radius": 0.035}, {"u": 0.75, "v": 0.70, "radius": 0.035}],
}


def png(path, pixels):
    def chunk(kind, data):
        return (
            struct.pack("!I", len(data))
            + kind
            + data
            + struct.pack("!I", zlib.crc32(kind + data) & 0xFFFFFFFF)
        )

    raw = b"".join(b"\0" + pixels[y * w * 4 : (y + 1) * w * 4] for y in range(h))
    path.write_bytes(
        b"\x89PNG\r\n\x1a\n"
        + chunk(b"IHDR", struct.pack("!IIBBBBB", w, h, 8, 6, 0, 0, 0))
        + chunk(b"IDAT", zlib.compress(raw))
        + chunk(b"IEND", b"")
    )


measure = {"kind": recipe["kind"], "blender": bpy.app.version_string, "recipe": recipe}
for painted in [False, True]:
    pixels = bytearray()
    for y in range(h):
        for x in range(w):
            u = (x + 0.5) / w
            v = 1 - (y + 0.5) / h
            c = [38, 104, 114, 255]
            if painted:
                a = max(0, min(1, (0.065 - abs(v - 0.5)) / 0.008))
                c = [round(c[i] * (1 - a) + [230, 127, 49, 255][i] * a) for i in range(4)]
                for center in [0.25, 0.75]:
                    dist = math.sqrt((u - center) ** 2 + (v - 0.70) ** 2)
                    a = max(0, min(1, (0.035 - dist) / 0.004))
                    c = [round(c[i] * (1 - a) + [230, 226, 202, 255][i] * a) for i in range(4)]
            pixels.extend(c)
    name = "texture-paint-after" if painted else "texture-paint-before"
    imagepath = out / (name + ".png")
    png(imagepath, pixels)
    bpy.ops.wm.read_factory_settings(use_empty=True)
    scene = bpy.context.scene
    scene.render.engine = "BLENDER_EEVEE"
    scene.render.resolution_x = 900
    scene.render.resolution_y = 720
    scene.render.resolution_percentage = 100
    scene.world = bpy.data.worlds.new("Studio world")
    scene.world.color = (0.12, 0.12, 0.12)
    scene.view_settings.view_transform = "AgX"
    coords = []
    for j in range(1, 32):
        theta = math.pi * j / 32
        for i in range(64):
            phi = 2 * math.pi * i / 64
            coords.append(
                (math.sin(theta) * math.cos(phi), math.sin(theta) * math.sin(phi), math.cos(theta))
            )
    north = len(coords)
    coords.append((0, 0, 1))
    south = len(coords)
    coords.append((0, 0, -1))
    faces = []
    face_uv = []
    for j in range(30):
        for i in range(64):
            nxt = (i + 1) % 64
            u = i / 64
            un = (i + 1) / 64
            v = 1 - (j + 1) / 32
            vn = 1 - (j + 2) / 32
            faces.append((j * 64 + i, (j + 1) * 64 + i, (j + 1) * 64 + nxt, j * 64 + nxt))
            face_uv.append([(u, v), (u, vn), (un, vn), (un, v)])
    for i in range(64):
        nxt = (i + 1) % 64
        u = i / 64
        un = (i + 1) / 64
        faces.append((north, i, nxt))
        face_uv.append([((u + un) / 2, 1), (u, 31 / 32), (un, 31 / 32)])
        faces.append((south, 30 * 64 + nxt, 30 * 64 + i))
        face_uv.append([((u + un) / 2, 0), (un, 1 / 32), (u, 1 / 32)])
    mesh = bpy.data.meshes.new("Explicit latitude longitude surface")
    mesh.from_pydata(coords, [], faces)
    mesh.update()
    obj = bpy.data.objects.new("Original UV painted prop", mesh)
    scene.collection.objects.link(obj)
    layer = mesh.uv_layers.new(name="Explicit UV map").data
    for face, uvs in zip(mesh.polygons, face_uv):
        for idx, coordinate in zip(face.loop_indices, uvs):
            layer[idx].uv = coordinate
    for p in obj.data.polygons:
        p.use_smooth = True
    image = bpy.data.images.load(str(imagepath))
    image.pack()
    mat = bpy.data.materials.new("UV color image")
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    tex = nodes.new("ShaderNodeTexImage")
    tex.image = image
    nodes.active = tex
    coord = nodes.new("ShaderNodeTexCoord")
    bs = nodes.get("Principled BSDF")
    bs.inputs["Roughness"].default_value = 0.6
    mat.node_tree.links.new(coord.outputs["UV"], tex.inputs["Vector"])
    mat.node_tree.links.new(tex.outputs["Color"], bs.inputs["Base Color"])
    obj.data.materials.append(mat)
    bpy.ops.object.camera_add(location=(4, -6, 3))
    cam = bpy.context.object
    cam.rotation_euler = (-cam.location).to_track_quat("-Z", "Y").to_euler()
    cam.data.type = "ORTHO"
    cam.data.ortho_scale = 3
    scene.camera = cam
    for loc, power, size in [((1, -3, 5), 900, 4), ((-3, 1, 2), 450, 3)]:
        bpy.ops.object.light_add(type="AREA", location=loc)
        light = bpy.context.object
        light.data.energy = power
        light.data.size = size
        light.rotation_euler = (-light.location).to_track_quat("-Z", "Y").to_euler()
    coords = [list(v.co) for v in obj.data.vertices]
    uv = [list(d.uv) for d in obj.data.uv_layers.active.data]
    measure[name] = {
        "vertices": len(coords),
        "vertexHash": hashlib.sha256(json.dumps(coords).encode()).hexdigest(),
        "uvHash": hashlib.sha256(json.dumps(uv).encode()).hexdigest(),
        "imageSHA256": hashlib.sha256(imagepath.read_bytes()).hexdigest(),
        "packed": bool(image.packed_file),
    }
    scene.render.image_settings.file_format = "PNG"
    scene.render.filepath = str(out / (name + "-render.png"))
    bpy.ops.wm.save_as_mainfile(filepath=str(out / (name + ".blend")))
    bpy.ops.render.render(write_still=True)
(out / "texture-paint-measurements.json").write_text(json.dumps(measure, indent=2) + "\n")