import bpy
import math
import mathutils
import json
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union

# ============================================================
# PATHS
# ============================================================

ROBOT_JSON_FILE = r"C:\Users\kech\SynologyDrive\2026-AppServer-AppRobot\appRobotRendering\robot.json"
OUTPUT_FILE = r"C:\Users\kech\SynologyDrive\2026-AppServer-AppRobot\appRobotRendering\render.png"
RENDER_WIDTH = 1200
RENDER_HEIGHT = 800

# ============================================================
# FALLBACK DEFAULT MATERIALS (placeholders)
# ============================================================

DEFAULT_MATERIALS = {
    "wood": {
        "baseColor": (0.72, 0.52, 0.33, 1.0),
        "roughness": 0.85,
        "metallic": 0.0,
    },
    "plaWhite": {
        "baseColor": (0.95, 0.95, 0.95, 1.0),
        "roughness": 0.45,
        "metallic": 0.0,
    },
    "steel": {
        "baseColor": (0.72, 0.72, 0.75, 1.0),
        "roughness": 0.25,
        "metallic": 1.0,
    },
    "powderCoatBlue": {
        "baseColor": (0.15, 0.25, 0.70, 1.0),
        "roughness": 0.55,
        "metallic": 0.0,
    },
    "marbleStone": {
        "baseColor": (0.85, 0.85, 0.87, 1.0),
        "roughness": 0.95,
        "metallic": 0.0,
    },
    "defaultPlastic": {
        "baseColor": (0.95, 0.95, 0.95, 1.0),
        "roughness": 0.40,
        "metallic": 0.0,
    },
}

STATE_KEYS = ["x", "y", "z", "a", "b", "c", "e"]

# ============================================================
# JSON LOADING
# ============================================================

robot: Dict[str, Any] = {}
if Path(ROBOT_JSON_FILE).exists():
    with open(ROBOT_JSON_FILE, "r", encoding="utf-8") as f:
        robot = json.load(f)
else:
    # Minimal fallback so the script can still run during development.
    robot = {
        "renderingInfo": {
            "cameraPosition": [-500, -1100, 900],
            "cameraTarget": [0, 0, 200],
            "cameraUpVector": [0, 0, 1],
            "lightPosition": [-1000, -1000, 2000],
            "lightTarget": [0, 0, 0],
            "lightUpVector": [0, 0, 1],
            "metric": "mm",
            "materials": {},
        },
        "defaultPosition": {k: 0 for k in STATE_KEYS},
        "recognized": {k: None for k in STATE_KEYS},
        "movements": {k: None for k in STATE_KEYS},
        "links": {},
    }

rendering_info = robot.get("renderingInfo", {})
metric = rendering_info.get("metric", "mm")
scale_factor = 0.001 if metric == "mm" else 1.0

# Merge current state from multiple places.
# Priority: movements -> recognized -> defaultPosition -> 0
state: Dict[str, float] = {k: 0.0 for k in STATE_KEYS}
for source_name in ("defaultPosition", "recognized", "movements"):
    source = robot.get(source_name, {}) or {}
    for k in STATE_KEYS:
        v = source.get(k, None)
        if v is not None:
            state[k] = float(v)

# ============================================================
# HELPERS
# ============================================================

def mm_to_m(value: float) -> float:
    return value * scale_factor


def resolve_scalar(value: Any, state_map: Dict[str, float]) -> float:
    """Resolve numbers or symbolic placeholders like 'x'/'a'."""
    if value is None:
        return 0.0
    if isinstance(value, (int, float)):
        return float(value)
    if isinstance(value, str):
        key = value.strip().lower()
        if key in state_map:
            return float(state_map[key])
        try:
            return float(key)
        except ValueError:
            return 0.0
    return 0.0


def resolve_vector(value: Any, state_map: Dict[str, float], default_len: int = 3) -> Tuple[float, ...]:
    if value is None:
        return tuple(0.0 for _ in range(default_len))
    if isinstance(value, (int, float, str)):
        return (resolve_scalar(value, state_map),)
    if isinstance(value, (list, tuple)):
        resolved = [resolve_scalar(v, state_map) for v in value]
        if len(resolved) < default_len:
            resolved.extend([0.0] * (default_len - len(resolved)))
        return tuple(resolved[:default_len])
    return tuple(0.0 for _ in range(default_len))


def resolve_vec3_m(value: Any, state_map: Dict[str, float]) -> Tuple[float, float, float]:
    x, y, z = resolve_vector(value, state_map, default_len=3)
    return mm_to_m(x), mm_to_m(y), mm_to_m(z)


def normalize_axis(axis: Iterable[Any]) -> mathutils.Vector:
    ax = mathutils.Vector((float(axis[0]), float(axis[1]), float(axis[2])))
    if ax.length == 0:
        return mathutils.Vector((1.0, 0.0, 0.0))
    return ax.normalized()


def create_or_get_material(name: str, fallback: str = "defaultPlastic") -> bpy.types.Material:
    if name in bpy.data.materials:
        return bpy.data.materials[name]

    info = (robot.get("renderingInfo", {}) or {}).get("materials", {}) or {}
    spec = None

    # Support both dict-style and old list-style material definitions.
    if isinstance(info, dict):
        spec = info.get(name)
    elif isinstance(info, list):
        for entry in info:
            if isinstance(entry, dict) and name in entry:
                spec = entry[name]
                break

    if not isinstance(spec, dict):
        spec = DEFAULT_MATERIALS.get(name, DEFAULT_MATERIALS[fallback])
    else:
        # Accept partial specs from JSON.
        base = DEFAULT_MATERIALS.get(name, DEFAULT_MATERIALS[fallback]).copy()
        if "baseColor" in spec:
            color = tuple(spec["baseColor"])
            base["baseColor"] = (*color[:3], 1.0) if len(color) == 3 else tuple(color[:4])
        if "roughness" in spec:
            base["roughness"] = float(spec["roughness"])
        if "metallic" in spec:
            base["metallic"] = float(spec["metallic"])
        spec = base

    mat = bpy.data.materials.new(name=name)
    mat.use_nodes = True
    bsdf = mat.node_tree.nodes.get("Principled BSDF")
    if bsdf is not None:
        bsdf.inputs["Base Color"].default_value = spec["baseColor"]
        bsdf.inputs["Roughness"].default_value = spec["roughness"]
        bsdf.inputs["Metallic"].default_value = spec["metallic"]
    return mat

def import_stl(filepath: str):

    filepath = str(Path(filepath).resolve())

    if not Path(filepath).exists():
        raise FileNotFoundError(
            f"STL file not found:\\n{filepath}"
        )

    before = set(bpy.data.objects)

    bpy.ops.wm.stl_import(filepath=filepath)

    after = [
        obj for obj in bpy.data.objects
        if obj not in before
    ]

    return after


def link_object(obj: bpy.types.Object):
    if obj.name not in bpy.context.collection.objects:
        bpy.context.collection.objects.link(obj)


def create_empty(name: str, location=(0, 0, 0), rotation=(0, 0, 0)) -> bpy.types.Object:
    empty = bpy.data.objects.new(name, None)
    bpy.context.collection.objects.link(empty)
    empty.location = location
    empty.rotation_euler = rotation
    return empty


def euler_deg_xyz(values: Any) -> Tuple[float, float, float]:
    x, y, z = resolve_vector(values, state, default_len=3)
    return math.radians(x), math.radians(y), math.radians(z)


def safe_parent(child: bpy.types.Object, parent: Optional[bpy.types.Object]):
    if parent is not None:
        child.parent = parent
        # Keep current world transform visually stable after parenting.
        child.matrix_parent_inverse = parent.matrix_world.inverted()


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

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

# ============================================================
# UNITS / WORLD
# ============================================================

scene = bpy.context.scene
scene.unit_settings.system = "METRIC"
scene.unit_settings.length_unit = "MILLIMETERS"
scene.unit_settings.scale_length = scale_factor

world = scene.world
if world is None:
    world = bpy.data.worlds.new("World")
    scene.world = world
world.use_nodes = True
bg = world.node_tree.nodes["Background"]
bg.inputs[0].default_value = (0.70, 0.85, 1.0, 1.0)  # light blue sky
bg.inputs[1].default_value = 0.2

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

scene.render.engine = "CYCLES"
scene.view_settings.exposure = -1.5
scene.cycles.samples = 64
scene.cycles.preview_samples = 32
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
scene.render.film_transparent = False

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

# 2m x 2m floor, centered at origin.
bpy.ops.mesh.primitive_plane_add(size=2.0, location=(0, 0, mm_to_m(-28.0)))
floor = bpy.context.active_object

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_node.inputs["Color1"].default_value = (0.82, 0.82, 0.82, 1.0)
checker_node.inputs["Color2"].default_value = (0.18, 0.18, 0.18, 1.0)

# 100mm checker squares across a 2m x 2m floor => 20 tiles each direction.
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)

# ============================================================
# 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 = resolve_vec3_m(rendering_info.get("cameraPosition", [-500, -1100, 900]), state)
cam_target = resolve_vec3_m(rendering_info.get("cameraTarget", [0, 0, 0]), state)

cam_obj.location = cam_pos
cam_data.lens = 50

cam_vec = mathutils.Vector(cam_target) - mathutils.Vector(cam_pos)
if cam_vec.length == 0:
    cam_vec = mathutils.Vector((1, 0, 0))
cam_obj.rotation_euler = cam_vec.to_track_quat("-Z", "Y").to_euler()
scene.camera = cam_obj

# ============================================================
# LIGHTS
# ============================================================

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)

sun_pos = resolve_vec3_m(rendering_info.get("lightPosition", [-1000, -1000, 2000]), state)
light_target = resolve_vec3_m(rendering_info.get("lightTarget", [0, 0, 0]), state)
sun_obj.location = sun_pos
light_vec = mathutils.Vector(light_target) - mathutils.Vector(sun_pos)
if light_vec.length == 0:
    light_vec = mathutils.Vector((1, 0, -1))
sun_obj.rotation_euler = light_vec.to_track_quat("-Z", "Y").to_euler()
sun_data.energy = 1.0

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 = (mm_to_m(-800), mm_to_m(-1200), mm_to_m(1500))
area_obj.rotation_euler = (math.radians(60), 0.0, math.radians(-20))
area_data.energy = 300
area_data.size = 2.0

# ============================================================
# ROBOT BUILDING
# ============================================================

links_def = robot.get("links")
if links_def is None:
    # Backward compatibility with older name.
    links_def = robot.get("ElementInfos", {})

created_nodes: Dict[str, bpy.types.Object] = {}

# Create all link containers first.
for link_name in links_def.keys():
    created_nodes[link_name] = create_empty(f"{link_name}_link")

# Parent/position link containers.
for link_name, link_info in links_def.items():
    parent_name = link_info.get("parent")
    parent_obj = created_nodes.get(parent_name) if parent_name else None

    link_obj = created_nodes[link_name]
    safe_parent(link_obj, parent_obj)

    # Static mounting transform relative to parent.
    # Keep the extra info, but rename it to mountRotation in your JSON.
    mount_pos = link_info.get("mountPosition", link_info.get("originInParentCoordinates", [0, 0, 0]))
    mount_rot = link_info.get("mountRotation", link_info.get("rotationInParentCoordinates", [0, 0, 0]))

    link_obj.location = resolve_vec3_m(mount_pos, state)
    link_obj.rotation_euler = euler_deg_xyz(mount_rot)

    # Joint transform (child-owned).
    joint = link_info.get("jointToParent") or link_info.get("joint")
    if isinstance(joint, dict):
        joint_origin = joint.get("origin", [0, 0, 0])
        joint_rot = joint.get("rotation", [0, 0, 0])
        joint_type = joint.get("type", "fixed")
        control_var = str(joint.get("variable", joint.get("control", ""))).lower()
        axis = joint.get("axis", [1, 0, 0])

        joint_offset = create_empty(f"{link_name}_joint")
        safe_parent(joint_offset, link_obj)
        joint_offset.location = resolve_vec3_m(joint_origin, state)
        joint_offset.rotation_euler = euler_deg_xyz(joint_rot)

        # Motion node under the joint offset.
        motion_node = create_empty(f"{link_name}_motion")
        safe_parent(motion_node, joint_offset)

        if joint_type == "linear":
            # Linear joint moves along its local axis by the control value.
            move_val_mm = state.get(control_var, 0.0) if control_var else 0.0
            axis_v = normalize_axis(axis)
            motion_node.location = axis_v * mm_to_m(move_val_mm)
        elif joint_type == "revolute":
            # Revolute joint rotates around its local axis by the control value.
            angle_deg = state.get(control_var, 0.0) if control_var else 0.0
            axis_v = normalize_axis(axis)
            # Convert axis-angle to Euler in local space by using rotation_difference.
            quat = mathutils.Quaternion(axis_v, math.radians(angle_deg))
            motion_node.rotation_euler = quat.to_euler()
        else:
            # fixed / unknown => no motion
            pass

        # The link container sits under the motion node.
        safe_parent(link_obj, motion_node)

# Import and attach all meshes for every link.
for link_name, link_info in links_def.items():
    link_obj = created_nodes[link_name]
    model_list = link_info.get("model", [])
    if not isinstance(model_list, list):
        model_list = []

    # Optional single-file shorthand.
    if "stlFile" in link_info:
        model_list = model_list + [{"stlFile": link_info["stlFile"]}]

    for idx, model_def in enumerate(model_list):
        stl_file = model_def.get("stlFile")
        if not stl_file:
            continue

        base_dir = Path(ROBOT_JSON_FILE).parent if Path(ROBOT_JSON_FILE).exists() else Path.cwd()
        stl_path = (base_dir / stl_file).resolve()
        if not stl_path.exists():
            # Try the file as given.
            stl_path = Path(stl_file).resolve()

        imported = import_stl(str(stl_path))

        # Create a mesh container for each imported STL so a link can have many surfaces.
        mesh_container = create_empty(f"{link_name}_mesh_{idx}")
        safe_parent(mesh_container, link_obj)

        origin_of_model = model_def.get("originOfModel", [0, 0, 0])
        rot_of_model = model_def.get("rotationOfModelDegree", [0, 0, 0])
        mesh_container.location = resolve_vec3_m(origin_of_model, state)
        mesh_container.rotation_euler = euler_deg_xyz(rot_of_model)

        material_name = model_def.get("material", "defaultPlastic")
        material = create_or_get_material(material_name)

        for obj in imported:
            if obj.type != "MESH":
                continue
            safe_parent(obj, mesh_container)
            # Keep STL imports at their own local origin; only scale to meters.
            obj.scale = (scale_factor, scale_factor, scale_factor)
            if len(obj.data.materials) == 0:
                obj.data.materials.append(material)
            else:
                obj.data.materials[0] = material


# ============================================================
# DEBUG WORLD AXES
# ============================================================

def create_axis_arrow(
    name,
    direction,
    color,
    length_mm=200,
    radius_mm=2,
    cone_radius_mm=5,
    cone_length_mm=20
):
    length = mm_to_m(length_mm)
    radius = mm_to_m(radius_mm)
    cone_radius = mm_to_m(cone_radius_mm)
    cone_length = mm_to_m(cone_length_mm)

    dir_vec = mathutils.Vector(direction).normalized()

    # --------------------------------------------------------
    # CYLINDER
    # --------------------------------------------------------

    bpy.ops.mesh.primitive_cylinder_add(
        radius=radius,
        depth=length - cone_length
    )

    cyl = bpy.context.active_object
    cyl.name = f"{name}_shaft"

    # Blender cylinder points along Z by default
    cyl.rotation_mode = 'QUATERNION'
    cyl.rotation_quaternion = (
        mathutils.Vector((0, 0, 1))
        .rotation_difference(dir_vec)
    )

    cyl.location = dir_vec * ((length - cone_length) * 0.5)

    # --------------------------------------------------------
    # CONE
    # --------------------------------------------------------

    bpy.ops.mesh.primitive_cone_add(
        radius1=cone_radius,
        depth=cone_length
    )

    cone = bpy.context.active_object
    cone.name = f"{name}_tip"

    cone.rotation_mode = 'QUATERNION'
    cone.rotation_quaternion = (
        mathutils.Vector((0, 0, 1))
        .rotation_difference(dir_vec)
    )

    cone.location = dir_vec * (length - cone_length * 0.5)

    # --------------------------------------------------------
    # MATERIAL
    # --------------------------------------------------------

    mat = bpy.data.materials.new(name=f"{name}_material")
    mat.use_nodes = True

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

    bsdf.inputs["Base Color"].default_value = (
        color[0],
        color[1],
        color[2],
        1.0
    )

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

    cyl.data.materials.append(mat)
    cone.data.materials.append(mat)


# ------------------------------------------------------------
# CREATE XYZ AXES
# ------------------------------------------------------------

# X = red
create_axis_arrow(
    "AxisX",
    (1, 0, 0),
    (1, 0, 0)
)

# Y = green
create_axis_arrow(
    "AxisY",
    (0, 1, 0),
    (0, 1, 0)
)

# Z = blue
create_axis_arrow(
    "AxisZ",
    (0, 0, 1),
    (0, 0, 1)
)


# ============================================================
# FINAL RENDER
# ============================================================

bpy.ops.render.render(write_still=True)
print("Finished rendering:", OUTPUT_FILE)