"""Build original scalloped tree meshes, editable in Blender and portable to Godot.

Run: Blender --background --python tools/art/build_painterly_trees.py
One crown surface and one joined trunk/branch surface; no textures or alpha cards.
All coordinates use Blender Z-up and glTF exports them as Godot Y-up.
"""

import json
import math
import random
from pathlib import Path

import bpy
from mathutils import Vector

ROOT = Path(__file__).resolve().parents[2]
SOURCE = ROOT / "art/painterly"
OUTPUT = ROOT / "assets/painterly"
SOURCE.mkdir(parents=True, exist_ok=True)
OUTPUT.mkdir(parents=True, exist_ok=True)
(SOURCE / ".gdignore").write_text("")
RNG = random.Random(2741)

bpy.ops.object.select_all(action="SELECT")
bpy.ops.object.delete(use_global=False)


def material(name):
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    shader = nodes.get("Principled BSDF")
    shader.inputs["Roughness"].default_value = 0.95
    shader.inputs["Specular IOR Level"].default_value = 0.0
    color = nodes.new("ShaderNodeVertexColor")
    color.layer_name = "Color"
    mat.node_tree.links.new(color.outputs["Color"], shader.inputs["Base Color"])
    return mat


def mesh_object(name, vertices, faces, colors, mat):
    mesh = bpy.data.meshes.new(name)
    mesh.from_pydata(vertices, [], faces)
    mesh.update()
    obj = bpy.data.objects.new(name, mesh)
    bpy.context.collection.objects.link(obj)
    attr = mesh.color_attributes.new(name="Color", type="FLOAT_COLOR", domain="POINT")
    for entry, color in zip(attr.data, colors):
        entry.color = (*color, 1.0)
    mesh.materials.append(mat)
    for polygon in mesh.polygons:
        polygon.use_smooth = True
    return obj


def source_tint(mat, tint):
    """Set an editing preview tint after portable glTF export."""
    nodes = mat.node_tree.nodes
    mix = nodes.new("ShaderNodeMixRGB")
    mix.blend_type = "MULTIPLY"
    mix.inputs[0].default_value = 1.0
    mix.inputs[2].default_value = (*tint, 1.0)
    mat.node_tree.links.new(nodes.get("Color Attribute").outputs["Color"], mix.inputs[1])
    mat.node_tree.links.new(mix.outputs[0], nodes.get("Principled BSDF").inputs["Base Color"])


# Connected, irregular crown with broad painted lobes and small leafy tips.
# Perturbing an envelope keeps the silhouette solid and removes internal overdraw.
bumps = []
for index in range(23):
    z = 1 - 2 * (index + 0.5) / 23
    angle = index * 2.399963 + 0.2
    direction = Vector((math.sqrt(1 - z * z) * math.cos(angle),
                        math.sqrt(1 - z * z) * math.sin(angle), z))
    bumps.append((direction, RNG.uniform(0.12, 0.25)))


def envelope(direction):
    ridges = sum(height * math.exp((direction.dot(axis) - 1) * 30)
                 for axis, height in bumps)
    fine = math.sin(direction.x * 22 + direction.z * 9) * math.sin(direction.y * 19) * 0.018
    radius = 0.84 + ridges + fine
    return Vector((direction.x * radius, direction.y * radius * 0.95,
                   direction.z * radius * 0.86))


def pigment(point):
    broad = math.sin(point.x * 8 + point.z * 2) * math.sin(point.y * 7 - point.z * 5)
    light = 0.82 + 0.11 * broad + 0.07 * max(0, point.z)
    return (light * (1.0 + 0.035 * broad), light, light * (0.95 - 0.025 * broad))


vertices, faces, colors = [], [], []
segments, rings = 32, 16
vertices.append(envelope(Vector((0, 0, 1))))
for ring in range(1, rings):
    phi = math.pi * ring / rings
    for segment in range(segments):
        theta = math.tau * segment / segments
        direction = Vector((math.sin(phi) * math.cos(theta), math.sin(phi) * math.sin(theta), math.cos(phi)))
        vertices.append(envelope(direction))
vertices.append(envelope(Vector((0, 0, -1))))
for segment in range(segments):
    faces.append((0, 1 + segment, 1 + (segment + 1) % segments))
for ring in range(rings - 2):
    for segment in range(segments):
        a = 1 + ring * segments + segment
        b = 1 + ring * segments + (segment + 1) % segments
        faces.append((a, a + segments, b + segments, b))
last = len(vertices) - 1
for segment in range(segments):
    faces.append((last, last - segments + (segment + 1) % segments, last - segments + segment))
colors = [pigment(point) for point in vertices]

# Bent solid leaf tips give a scalloped edge at close range. They are tiny opaque
# wedges with both sides, so they do not require cull-disabled or alpha materials.
for axis, _height in bumps:
    for leaf in range(2):
        tangent = axis.cross(Vector((0, 0, 1)))
        if tangent.length < 0.01:
            tangent = Vector((1, 0, 0))
        tangent.normalize()
        sideways = axis.cross(tangent).normalized()
        direction = (axis + tangent * (leaf - 0.5) * 0.13).normalized()
        base = envelope(direction) * 0.98
        length = RNG.uniform(0.11, 0.17)
        width = RNG.uniform(0.04, 0.065)
        middle = base + direction * length * 0.55
        start = len(vertices)
        vertices.extend([base, middle + tangent * width, middle + sideways * width * 0.28,
                         middle - tangent * width, base + direction * length])
        faces.extend([tuple(start + j for j in indices) for indices in
                      ((0, 1, 2), (0, 2, 3), (1, 4, 2), (2, 4, 3), (0, 3, 1), (1, 3, 4))])
        colors.extend([pigment(point) for point in vertices[-5:]])
canopy = mesh_object("Canopy", vertices, faces, colors, material("CanopyPigment"))

# A crooked trunk, joined forks and roots use one surface. Colors are linear,
# with grey warm bark, moss at the base and subtle broad plane variation.
vertices, faces, colors = [], [], []


def tube(points, radii, sides=9):
    start = len(vertices)
    for index, raw in enumerate(points):
        point = Vector(raw)
        tangent = Vector(points[min(index + 1, len(points) - 1)]) - Vector(points[max(0, index - 1)])
        tangent.normalize()
        side = tangent.cross(Vector((0, 1, 0))).normalized()
        other = tangent.cross(side).normalized()
        for segment in range(sides):
            angle = math.tau * segment / sides
            offset = (math.cos(angle) * side + math.sin(angle) * other) * radii[index]
            vertices.append(point + offset)
            shade = 0.90 + math.sin(angle * 3 + 0.6) * 0.12
            moss = max(0, 1 - point.z / 0.55) * (0.5 + 0.5 * math.cos(angle))
            colors.append(((0.20 - moss * 0.035) * shade, (0.13 + moss * 0.01) * shade,
                           (0.068 - moss * 0.005) * shade))
        if index:
            for segment in range(sides):
                a = start + (index - 1) * sides + segment
                b = start + (index - 1) * sides + (segment + 1) % sides
                faces.append((a, b, b + sides, a + sides))
    faces.append(tuple(reversed(range(start, start + sides))))
    faces.append(tuple(range(len(vertices) - sides, len(vertices))))


tube([(0, 0, 0), (-0.07, 0.025, 0.22), (0.02, 0.0, 1.0), (-0.10, 0.025, 1.85),
      (0.03, 0.02, 2.62), (0.22, -0.04, 3.48)], [0.38, 0.27, 0.23, 0.20, 0.14, 0.025])
tube([(-0.06, 0, 1.70), (-0.36, 0.04, 2.35), (-0.84, 0.03, 2.83), (-1.05, 0.09, 3.23)],
     [0.18, 0.13, 0.075, 0.012], 7)
tube([(0.00, 0, 2.10), (0.39, 0.12, 2.48), (0.81, 0.18, 2.94), (1.00, 0.25, 3.29)],
     [0.14, 0.10, 0.055, 0.012], 7)
tube([(-0.06, 0.025, 2.34), (-0.22, -0.38, 2.83), (-0.35, -0.62, 3.26)],
     [0.10, 0.058, 0.012], 7)
for index in range(5):
    angle = index * math.tau / 5 + 0.2
    tube([(math.cos(angle) * 0.13, math.sin(angle) * 0.13, 0.30),
          (math.cos(angle) * 0.36, math.sin(angle) * 0.36, 0.08),
          (math.cos(angle + 0.12) * 0.66, math.sin(angle + 0.12) * 0.66, -0.015)],
         [0.12, 0.085, 0.012], 6)
trunk = mesh_object("Trunk", vertices, faces, colors, material("BarkPigment"))

# Export object origins at zero; runtime scales/repositions four shared crowns.
bpy.ops.object.select_all(action="SELECT")
bpy.ops.export_scene.gltf(filepath=str(OUTPUT / "painterly_tree.glb"), export_format="GLB",
                         use_selection=True, export_yup=True, export_animations=False,
                         export_vertex_color="ACTIVE", export_all_vertex_colors=False,
                         export_materials="EXPORT", export_extras=False, export_cameras=False)

# Display a complete tree in the editable source without baking placement into GLB.
source_tint(canopy.data.materials[0], (0.22, 0.42, 0.10))
canopy.location.z = 3.8
canopy.scale = (1.8, 1.55, 1.4)
for location, scale in (((-0.85, 0.05, 3.55), (1.28, 1.22, 1.15)),
                        ((0.87, -0.08, 3.67), (1.27, 1.15, 1.1)),
                        ((0.05, 0, 4.46), (1.20, 1.08, 1.0))):
    crown = canopy.copy()
    crown.data = canopy.data
    bpy.context.collection.objects.link(crown)
    crown.location = location
    crown.scale = scale
for screen in bpy.data.screens:
    for area in screen.areas:
        if area.type == "VIEW_3D":
            area.spaces.active.shading.type = "MATERIAL"
bpy.context.preferences.filepaths.save_version = 0
bpy.ops.wm.save_as_mainfile(filepath=str(SOURCE / "painterly_tree.blend"))

budget = {}
for obj in (canopy, trunk):
    obj.data.calc_loop_triangles()
    budget[obj.name] = {"triangles": len(obj.data.loop_triangles), "surfaces": 1,
                        "vertices": len(obj.data.vertices)}
budget["runtime_tree"] = {"triangles": budget["Canopy"]["triangles"] * 4 + budget["Trunk"]["triangles"],
                          "surfaces": 5, "shared_meshes": 2}
budget["glb_bytes"] = (OUTPUT / "painterly_tree.glb").stat().st_size
(OUTPUT / "tree_mesh_budget.json").write_text(json.dumps(budget, indent=2) + "\n")
print("PAINTERLY TREE BUDGET", json.dumps(budget))

# Narrow evergreen boughs, scalloped skirts and an uneven leader. One opaque
# surface is joined across eight whorls, preserving a restrained draw budget.
bpy.ops.object.select_all(action="SELECT")
bpy.ops.object.delete(use_global=False)
vertices, faces, colors = [], [], []
segments = 24
for tier in range(8):
    height = 0.64 + tier * 0.27
    radius = 0.62 * (1 - tier / 8.8)
    start = len(vertices)
    angle_offset = tier * 0.61
    lean_x = math.sin(tier * 0.8) * 0.035
    lean_y = math.cos(tier * 1.1) * 0.025
    for ring, (radial, lift) in enumerate(((0.78, 0.04), (1.0, 0.08), (0.63, 0.25), (0.12, 0.57))):
        for segment in range(segments):
            angle = math.tau * segment / segments + angle_offset
            tip = (1 + math.cos(angle * 8 + tier)) * 0.5
            organic = math.sin(angle * 3 + tier) * 0.08 + math.sin(angle * 5 - tier * 2) * 0.035
            r = radius * radial * (0.84 + tip * 0.16 + organic)
            z = height + lift - tip * (0.085 if ring < 2 else 0.022)
            vertices.append((lean_x + math.cos(angle) * r, lean_y + math.sin(angle) * r, z))
            pigment = 0.72 + tier * 0.015 + ring * 0.05 + tip * 0.045
            colors.append((pigment * 0.94, pigment, pigment * 0.93))
            if ring:
                a = start + (ring - 1) * segments + segment
                b = start + (ring - 1) * segments + (segment + 1) % segments
                faces.append((a, b, b + segments, a + segments))
    faces.append(tuple(reversed(range(start, start + segments))))
    faces.append(tuple(range(len(vertices) - segments, len(vertices))))
conifer = mesh_object("ConiferCanopy", vertices, faces, colors, material("EvergreenPigment"))
vertices, faces, colors = [], [], []
tube([(0, 0, 0), (0.015, -0.01, 0.48), (-0.02, 0, 1.5), (0, 0, 2.58), (0.01, 0, 2.84)],
     [0.10, 0.07, 0.045, 0.02, 0.004], 8)
conifer_trunk = mesh_object("ConiferTrunk", vertices, faces, colors, material("EvergreenBark"))
bpy.ops.object.select_all(action="SELECT")
bpy.ops.export_scene.gltf(filepath=str(OUTPUT / "painterly_conifer.glb"), export_format="GLB",
                         use_selection=True, export_yup=True, export_animations=False,
                         export_vertex_color="ACTIVE", export_all_vertex_colors=False,
                         export_materials="EXPORT", export_extras=False, export_cameras=False)
source_tint(conifer.data.materials[0], (0.10, 0.28, 0.17))
bpy.ops.wm.save_as_mainfile(filepath=str(SOURCE / "painterly_conifer.blend"))
for obj in (conifer, conifer_trunk):
    obj.data.calc_loop_triangles()
budget["conifer"] = {"triangles": sum(len(o.data.loop_triangles) for o in (conifer, conifer_trunk)),
                      "surfaces": 2, "height_m": 3.10,
                      "glb_bytes": (OUTPUT / "painterly_conifer.glb").stat().st_size}
(OUTPUT / "tree_mesh_budget.json").write_text(json.dumps(budget, indent=2) + "\n")
print("CONIFER BUDGET", json.dumps(budget["conifer"]))