"""Create original native staging references without any generative model execution."""

import argparse
import json
import math
import sys
from pathlib import Path

import bpy
from mathutils import Vector

p = argparse.ArgumentParser()
p.add_argument("directory", type=Path)
a = p.parse_args(sys.argv[sys.argv.index("--") + 1 :])
r = a.directory
if r.exists():
    raise SystemExit("Choose a new output directory")
r.mkdir(parents=True)
bpy.ops.object.select_all(action="SELECT")
bpy.ops.object.delete(use_global=False)
s = bpy.context.scene
s.render.engine = "CYCLES"
s.cycles.samples = 24
s.cycles.use_denoising = True
s.render.resolution_percentage = 100
s.render.image_settings.file_format = "PNG"
s.render.pixel_aspect_x = 1
s.render.pixel_aspect_y = 1
s.world.color = (0.035, 0.035, 0.035)
s.view_settings.view_transform = "AgX"


def mat(name, color, metal=0):
    m = bpy.data.materials.new(name)
    m.diffuse_color = (*color, 1)
    m.use_nodes = True
    n = m.node_tree.nodes.get("Principled BSDF")
    n.inputs["Base Color"].default_value = (*color, 1)
    n.inputs["Metallic"].default_value = metal
    n.inputs["Roughness"].default_value = 0.32
    return m


teal = mat("Original teal bottle", (0.045, 0.31, 0.34), 0.15)
coral = mat("Original coral band", (0.85, 0.19, 0.08))
cream = mat("Warm ivory", (0.8, 0.7, 0.51))
floor = mat("Deep blue floor", (0.018, 0.035, 0.07))


def cyl(name, radius, depth, z, material):
    bpy.ops.mesh.primitive_cylinder_add(vertices=96, radius=radius, depth=depth, location=(0, 0, z))
    o = bpy.context.object
    o.name = name
    o.data.materials.append(material)
    b = o.modifiers.new("Rounded edges", "BEVEL")
    b.width = 0.055
    b.segments = 3
    o.modifiers.new("Surface normals", "WEIGHTED_NORMAL")
    return o


cyl("Bottle body", 0.53, 1.65, 1.0, teal)
cyl("Bottle band", 0.539, 0.22, 0.95, coral)
cyl("Bottle cap", 0.42, 0.23, 1.94, cream)
bpy.ops.mesh.primitive_plane_add(size=200)
bpy.context.object.data.materials.append(floor)
bpy.ops.object.camera_add(location=(6, -9, 5))
cam = bpy.context.object
cam.rotation_euler = (Vector((0, 0, 1.0)) - cam.location).to_track_quat("-Z", "Y").to_euler()
cam.data.lens = 52
s.camera = cam
for loc, power, size, col in [
    ((1, -3, 6), 850, 4, (0.75, 0.88, 1)),
    ((-4, 0, 4), 1100, 3, (1, 0.6, 0.25)),
]:
    bpy.ops.object.light_add(type="AREA", location=loc)
    o = bpy.context.object
    o.data.energy = power
    o.data.shape = "DISK"
    o.data.size = size
    o.data.color = col
    o.rotation_euler = (Vector((0, 0, 1)) - o.location).to_track_quat("-Z", "Y").to_euler()
s.render.resolution_x = 1000
s.render.resolution_y = 667
s.render.filepath = str(r / "hidream-native-reference.png")
bpy.ops.wm.save_as_mainfile(filepath=str(r / "light-hidream-reference.blend"))
bpy.ops.render.render(write_still=True)
curve = bpy.data.curves.new("Authored luminous arch", "CURVE")
curve.dimensions = "3D"
curve.bevel_depth = 0.025
curve.bevel_resolution = 3
sp = curve.splines.new("POLY")
sp.points.add(120)
for i, pt in enumerate(sp.points):
    t = i / 120
    pt.co = (-2.6 + 5.2 * t, 0.65, 0.65 + 2.1 * math.sin(math.pi * t), 1)
o = bpy.data.objects.new("Illustrative light path not a captured exposure", curve)
bpy.context.collection.objects.link(o)
em = bpy.data.materials.new("Warm illustrative trail")
em.use_nodes = True
nt = em.node_tree
nt.nodes.clear()
out = nt.nodes.new("ShaderNodeOutputMaterial")
e = nt.nodes.new("ShaderNodeEmission")
e.inputs["Color"].default_value = (1, 0.48, 0.1, 1)
e.inputs["Strength"].default_value = 4
nt.links.new(e.outputs[0], out.inputs["Surface"])
curve.materials.append(em)
s.render.resolution_x = 1600
s.render.resolution_y = 900
s.render.filepath = str(r / "light-painting-native-reference.png")
bpy.ops.wm.save_as_mainfile(filepath=str(r / "light-painting-staging.blend"))
bpy.ops.render.render(write_still=True)
(r / "light-hidream-native-record.json").write_text(
    json.dumps(
        {
            "Blender": bpy.app.version_string,
            "scope": "Original native staging renders only not photographed exposures or model outputs",
            "reference": [1000, 667],
            "trail": [1600, 900],
            "trailVertices": 121,
            "cameraMatrix": [list(x) for x in cam.matrix_world],
            "squarePixels": True,
        },
        indent=2,
    )
)
