import bpy
import mathutils

# ============================================================
# CONFIG
# ============================================================

robot = {
    "vision_config":{
        "MarkerType":"DICT_4X4_250",
        "MarkerSize":0.025
    },
    "renderingInfo":{
        "cameraPosition":[-400, -700, 300],
        "cameraTarget":["x", 0, 0],
        "cameraUpVector":[0, 0, 1],
        "lightPosition":[-500, -500, 500],
        "lightTarget":[0, 0, 0],
        "lightUpVector":[0, 0, 1],
        "metric": "mm",
        "materials":{
            "wood":{
                "baseColor":[0.72,0.52,0.33],
                "roughness":0.8,
                "metallic":0.0
            },

            "plaWhite":{
                "baseColor":[0.95,0.95,0.95],
                "roughness":0.45,
                "metallic":0.0
            },

            "steel":{
                "baseColor":[0.7,0.7,0.72],
                "roughness":0.25,
                "metallic":1.0
            },

            "powderCoatBlue":{
                "baseColor":[0.15,0.25,0.7],
                "roughness":0.55,
                "metallic":0.0
            },

            "marbleStone":{
                "baseColor":[0.85,0.85,0.87],
                "roughness":0.9,
                "metallic":0.0
            }
        }
    },
    "recognized":{"x":None, "y":None, "z": None, "a":None, "b":None, "c":None, "e": None},
    "movements":{"x":None, "y":None, "z": None, "a":None, "b":None, "c":None, "e": None},
    "elements":{
        "Board":{"type":"static", "parent":None, "size":[1000, 200, 25], 
                 "model":[{
                        "stlFile":"Board.stl",
                        "originOfModel":[0,0,0],
                        "rotationOfModelDegree":[0,0,90],
                        "material":"wood"},
                    {
                        "stlFile":"BoardRail.stl",
                        "originOfModel":[0,0,0],
                        "rotationOfModelDegree":[0,0,90],
                        "material":"steel"}
                    ]},
        "Base":{"type":"rigid", "parent":"Board", "size":[150, 200, 150],  "rotationInParentCoordinates":[0, 0, 0], 
            "model":[{  "stlFile":"Base.stl"
                    }],
            "jointToParent":{"name":"Slider", "type":"linear", "axis":[1,0,0], "origin":[0, 0, 0], 
                     "rotation":[0, 0, 0],
                     "variable":"x"}},
        "Arm1":{"type":"rigid", "parent":"Base", "size":[70, 250, 70], 
            "model":[{  "stlFile":"Holm.stl", 
                        "originOfModel":[0,0,0],
                        "rotationOfModelDegree":[0,0,0],
                        "material":"powderCoatBlue"
                    }]
        }
    }
}



LOCAL_PATH = r"C:\Users\kech\SynologyDrive\2026-AppServer-AppRobot\appRobotRendering\\"
MODEL_FILE = LOCAL_PATH + r"surfaces\BoardRail.stl"

OUTPUT_FILE = r"C:\Users\kech\SynologyDrive\2026-AppServer-AppRobot\appRobotRendering\render.png"

RENDER_WIDTH = 2000
RENDER_HEIGHT = 1000

# ============================================================
# CLEAN SCENE
# ============================================================

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

# ============================================================
# UNITS
# ============================================================

scene = bpy.context.scene

scene.unit_settings.system = 'METRIC'
scene.unit_settings.length_unit = 'MILLIMETERS'

metric = robot["renderingInfo"]["metric"]

# IMPORTANT:
# Blender internally uses meters.
# Your STL is already in millimeters.
# We therefore scale mm -> meters.
scale_factor = 0.001 if metric == "mm" else 1.0

# ============================================================
# IMPORT STL
# ============================================================

def import_stl(filepath):
    try:
        bpy.ops.wm.stl_import(filepath=filepath)
    except:
        bpy.ops.import_mesh.stl(filepath=filepath)

import_stl(MODEL_FILE)

imported_objects = bpy.context.selected_objects

# Apply scale
for obj in imported_objects:
    obj.scale = (scale_factor, scale_factor, scale_factor)

# ============================================================
# CENTER OBJECT
# ============================================================

bpy.ops.object.select_all(action='DESELECT')

for obj in imported_objects:
    obj.select_set(True)

bpy.context.view_layer.objects.active = imported_objects[0]

# Move origin to geometry center
bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='BOUNDS')

# Move object to world center
for obj in imported_objects:
    obj.location = (0, 0, 0)

# ============================================================
# WHITE PLASTIC MATERIAL
# ============================================================

mat = bpy.data.materials.new(name="WhitePlastic")
mat.use_nodes = True

bsdf = mat.node_tree.nodes["Principled BSDF"]

bsdf.inputs["Base Color"].default_value = (
    0.95, 0.95, 0.95, 1.0
)

bsdf.inputs["Roughness"].default_value = 0.4
bsdf.inputs["Metallic"].default_value = 0.0

for obj in imported_objects:
    if obj.type == 'MESH':

        if len(obj.data.materials) == 0:
            obj.data.materials.append(mat)
        else:
            obj.data.materials[0] = mat

# ============================================================
# CAMERA
# ============================================================

cam_data = bpy.data.cameras.new("Camera")
cam_obj = bpy.data.objects.new("Camera", cam_data)

bpy.context.collection.objects.link(cam_obj)

cam_pos = robot["renderingInfo"]["cameraPosition"]
cam_target = robot["renderingInfo"]["cameraTarget"]

# Convert mm -> meters
cam_obj.location = (
    cam_pos[0] * scale_factor,
    cam_pos[1] * scale_factor,
    cam_pos[2] * scale_factor
)

target = mathutils.Vector((
    cam_target[0] * scale_factor,
    cam_target[1] * scale_factor,
    cam_target[2] * scale_factor
))

direction = target - cam_obj.location

cam_obj.rotation_euler = (
    direction.to_track_quat('-Z', 'Y').to_euler()
)

cam_data.lens = 50

scene.camera = cam_obj

# ============================================================
# SUN LIGHT
# ============================================================

sun_data = bpy.data.lights.new(name="Sun", type='SUN')
sun_obj = bpy.data.objects.new(name="Sun", object_data=sun_data)

bpy.context.collection.objects.link(sun_obj)

light_pos = robot["renderingInfo"]["lightPosition"]
light_target = robot["renderingInfo"]["lightTarget"]

sun_obj.location = (
    light_pos[0] * scale_factor,
    light_pos[1] * scale_factor,
    light_pos[2] * scale_factor
)

light_target_vec = mathutils.Vector((
    light_target[0] * scale_factor,
    light_target[1] * scale_factor,
    light_target[2] * scale_factor
))

light_direction = light_target_vec - sun_obj.location

sun_obj.rotation_euler = (
    light_direction.to_track_quat('-Z', 'Y').to_euler()
)

sun_data.energy = 3.0

# ============================================================
# ADDITIONAL AREA LIGHT
# ============================================================

area_data = bpy.data.lights.new(name="AreaLight", type='AREA')
area_obj = bpy.data.objects.new(
    name="AreaLight",
    object_data=area_data
)

bpy.context.collection.objects.link(area_obj)

area_obj.location = (0, -1.2, 15)

area_data.energy = 5000
area_data.size = 2.0

# ============================================================
# CHECKERBOARD FLOOR
# ============================================================

# 2m x 2m floor
bpy.ops.mesh.primitive_plane_add(size=2, location=(0, 0, -27))

floor = bpy.context.active_object

# Create checker material
checker_mat = bpy.data.materials.new(name="Checkerboard")
checker_mat.use_nodes = True

nodes = checker_mat.node_tree.nodes
links = checker_mat.node_tree.links

nodes.clear()

output_node = nodes.new(type='ShaderNodeOutputMaterial')
bsdf_node = nodes.new(type='ShaderNodeBsdfPrincipled')
checker_node = nodes.new(type='ShaderNodeTexChecker')
mapping_node = nodes.new(type='ShaderNodeMapping')
texcoord_node = nodes.new(type='ShaderNodeTexCoord')

# Checker colors
checker_node.inputs["Color1"].default_value = (
    0.8, 0.8, 0.8, 1.0
)

checker_node.inputs["Color2"].default_value = (
    0.2, 0.2, 0.2, 1.0
)

# 100mm tiles
# floor is 2m -> 20 tiles
mapping_node.inputs["Scale"].default_value = (
    20.0, 20.0, 20.0
)

links.new(
    texcoord_node.outputs["UV"],
    mapping_node.inputs["Vector"]
)

links.new(
    mapping_node.outputs["Vector"],
    checker_node.inputs["Vector"]
)

links.new(
    checker_node.outputs["Color"],
    bsdf_node.inputs["Base Color"]
)

links.new(
    bsdf_node.outputs["BSDF"],
    output_node.inputs["Surface"]
)

floor.data.materials.append(checker_mat)

# ============================================================
# SKY BACKGROUND
# ============================================================

world = scene.world
world.use_nodes = True

bg = world.node_tree.nodes["Background"]

# Light blue sky
bg.inputs[0].default_value = (
    0.70, 0.85, 1.0, 1.0
)

bg.inputs[1].default_value = 0.8

# ============================================================
# RENDER SETTINGS
# ============================================================

scene.render.engine = 'CYCLES'

scene.cycles.samples = 128

scene.render.resolution_x = RENDER_WIDTH
scene.render.resolution_y = RENDER_HEIGHT
scene.render.resolution_percentage = 100

scene.render.image_settings.file_format = 'PNG'

scene.render.filepath = OUTPUT_FILE

# Slightly nicer shadows
scene.cycles.preview_samples = 32

# ============================================================
# RENDER
# ============================================================

bpy.ops.render.render(write_still=True)

print("Finished rendering:")
print(OUTPUT_FILE)