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

# ============================================================
# 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

# ============================================================
# DEFAULT MATERIALS
# ============================================================

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},
    "defaultPlastic": {"baseColor": (0.95, 0.95, 0.95, 1.0), "roughness": 0.40, "metallic": 0.0},
    "skeletonRed": {"baseColor": (0.85, 0.20, 0.20, 1.0), "roughness": 0.35, "metallic": 0.0},
    "markerBlack": {"baseColor": (0.04, 0.04, 0.04, 1.0), "roughness": 0.80, "metallic": 0.0},
}

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

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

with open(ROBOT_JSON_FILE, "r", encoding="utf-8") as f:
    robot: Dict[str, Any] = json.load(f)

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

def as_bool(value: Any, default: bool = False) -> bool:
    if value is None:
        return default
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        return value.strip().lower() in ("1", "true", "yes", "on")
    return bool(value)

show_skeleton = as_bool(rendering_info.get("showSkeleton", False))
show_markers = as_bool(rendering_info.get("showMarkers", False))

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:
            try:
                state[k] = float(v)
            except Exception:
                pass

links_def = robot.get("links", {})
if not isinstance(links_def, dict):
    raise ValueError("robot.json must contain a top-level 'links' object")

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

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

def resolve_scalar(value: Any, state_map: Dict[str, float]) -> float:
    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]:
    vec = list(resolve_vector(value, state_map, default_len=3))
    while len(vec) < 3:
        vec.append(0.0)
    x, y, z = vec[: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])))
    return ax.normalized() if ax.length > 0 else mathutils.Vector((1.0, 0.0, 0.0))

def euler_deg_xyz(values: Any) -> Tuple[float, float, float]:
    vec = list(resolve_vector(values, state, default_len=3))
    while len(vec) < 3:
        vec.append(0.0)
    return math.radians(vec[0]), math.radians(vec[1]), math.radians(vec[2])

def create_or_get_material(name: str, fallback: str = "defaultPlastic") -> bpy.types.Material:
    info = rendering_info.get("materials", {}) or {}
    spec = None

    if isinstance(info, dict):
        spec = info.get(name)

    if isinstance(spec, dict):
        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
    else:
        spec = DEFAULT_MATERIALS.get(name, DEFAULT_MATERIALS[fallback])

    if name in bpy.data.materials:
        mat = bpy.data.materials[name]
    else:
        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) -> List[bpy.types.Object]:
    path = Path(filepath).resolve()
    if not path.exists():
        raise FileNotFoundError(f"STL file not found:\n{path}")

    before = set(bpy.data.objects)
    bpy.ops.wm.stl_import(filepath=str(path))
    after = [obj for obj in bpy.data.objects if obj not in before]
    return after

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

def safe_parent(child: bpy.types.Object, parent: Optional[bpy.types.Object]):
    if parent is not None:
        child.parent = parent
        child.matrix_parent_inverse = parent.matrix_world.inverted()

def create_material_segment(name: str, color: Tuple[float, float, float], roughness: float = 0.35) -> bpy.types.Material:
    if name in bpy.data.materials:
        mat = bpy.data.materials[name]
    else:
        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 = (color[0], color[1], color[2], 1.0)
        bsdf.inputs["Roughness"].default_value = roughness
        bsdf.inputs["Metallic"].default_value = 0.0
    return mat

def create_cylinder_between(
    name: str,
    p1_local: Tuple[float, float, float],
    p2_local: Tuple[float, float, float],
    radius_m: float,
    parent: bpy.types.Object,
    material: bpy.types.Material
) -> bpy.types.Object:
    v1 = mathutils.Vector(p1_local)
    v2 = mathutils.Vector(p2_local)
    delta = v2 - v1
    length = delta.length
    if length <= 1e-9:
        length = 1e-6
        delta = mathutils.Vector((0.0, 0.0, 1e-6))

    bpy.ops.mesh.primitive_cylinder_add(radius=radius_m, depth=length)
    obj = bpy.context.active_object
    obj.name = name
    safe_parent(obj, parent)
    obj.location = (v1 + v2) * 0.5
    obj.rotation_mode = "QUATERNION"
    obj.rotation_quaternion = mathutils.Vector((0, 0, 1)).rotation_difference(delta.normalized())
    if len(obj.data.materials) == 0:
        obj.data.materials.append(material)
    else:
        obj.data.materials[0] = material
    return obj

def derive_default_skeleton_from_size(size_mm: List[float]) -> Dict[str, Any]:
    sx, sy, sz = (float(size_mm[0]), float(size_mm[1]), float(size_mm[2]))
    ax = max((abs(sx), 0), (abs(sy), 1), (abs(sz), 2), key=lambda x: x[0])[1]

    if ax == 0:
        return {"from": [0, sy * 0.5, sz * 0.5], "to": [sx, sy * 0.5, sz * 0.5]}
    if ax == 1:
        return {"from": [sx * 0.5, 0, sz * 0.5], "to": [sx * 0.5, sy, sz * 0.5]}
    return {"from": [sx * 0.5, sy * 0.5, 0], "to": [sx * 0.5, sy * 0.5, sz]}

def resolve_stl_path(stl_file: str) -> Path:
    base_dir = Path(ROBOT_JSON_FILE).parent
    candidates = [
        base_dir / stl_file,
        base_dir / "surfaces" / stl_file,
        Path(stl_file),
    ]
    for c in candidates:
        p = c.resolve()
        if p.exists():
            return p
    raise FileNotFoundError(
        "STL file not found in any expected location:\n" +
        "\n".join(str(c.resolve()) for c in candidates)
    )

# ============================================================
# SCENE RESET
# ============================================================

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

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

# ============================================================
# WORLD / RENDER SETTINGS
# ============================================================

world = scene.world or bpy.data.worlds.new("World")
scene.world = world
world.use_nodes = True
bg = world.node_tree.nodes["Background"]
bg.inputs[0].default_value = tuple(rendering_info.get("backgroundColor", [0.70, 0.85, 1.0])) + (1.0,)
bg.inputs[1].default_value = float(rendering_info.get("backgroundStrength", 0.20))

scene.render.engine = "CYCLES"
scene.view_settings.exposure = float(rendering_info.get("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
# ============================================================

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)
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", [-400, -700, 300]), 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", [-500, -500, 500]), 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 = float(rendering_info.get("sunEnergy", 0.35))

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 = float(rendering_info.get("areaEnergy", 120))
area_data.size = 2.0

# ============================================================
# ROBOT HIERARCHY
# ============================================================

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

for link_name in links_def.keys():
    link_frames[link_name] = create_empty(f"{link_name}_frame")

for link_name, link_info in links_def.items():
    parent_name = link_info.get("parent")
    parent_frame = link_frames.get(parent_name) if parent_name else None
    size_mm = link_info.get("size", [100, 100, 100])

    mount = create_empty(f"{link_name}_mount")
    safe_parent(mount, parent_frame)
    mount.location = resolve_vec3_m(link_info.get("mountPosition", [0, 0, 0]), state)
    mount.rotation_euler = euler_deg_xyz(link_info.get("mountRotation", [0, 0, 0]))

    joint_info = link_info.get("jointToParent", {}) or {}
    joint = create_empty(f"{link_name}_joint")
    safe_parent(joint, mount)
    joint.location = resolve_vec3_m(joint_info.get("origin", [0, 0, 0]), state)
    joint.rotation_euler = euler_deg_xyz(joint_info.get("rotation", [0, 0, 0]))

    motion = create_empty(f"{link_name}_motion")
    safe_parent(motion, joint)

    joint_type = str(joint_info.get("type", "fixed")).lower()
    control_var = str(joint_info.get("variable", joint_info.get("control", ""))).lower()
    axis = joint_info.get("axis", [1, 0, 0])

    if joint_type == "linear":
        value_mm = state.get(control_var, 0.0) if control_var else 0.0
        motion.location = normalize_axis(axis) * mm_to_m(value_mm)
    elif joint_type == "revolute":
        value_deg = state.get(control_var, 0.0) if control_var else 0.0
        motion.rotation_mode = "QUATERNION"
        motion.rotation_quaternion = mathutils.Quaternion(normalize_axis(axis), math.radians(value_deg))

    link_frame = link_frames[link_name]
    safe_parent(link_frame, motion)

    # --------------------------------------------------------
    # VISUAL MESHES
    # --------------------------------------------------------

    visual_root = create_empty(f"{link_name}_visual")
    safe_parent(visual_root, link_frame)

    model_list = link_info.get("model", [])
    if not isinstance(model_list, list):
        model_list = []

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

        stl_path = resolve_stl_path(stl_file)
        imported = import_stl(str(stl_path))

        model_node = create_empty(f"{link_name}_model_{idx}")
        safe_parent(model_node, visual_root)
        model_node.location = resolve_vec3_m(model_def.get("originOfModel", [0, 0, 0]), state)
        model_node.rotation_euler = euler_deg_xyz(model_def.get("rotationOfModelDegree", [0, 0, 0]))

        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, model_node)
            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

    # --------------------------------------------------------
    # SKELETON DEBUG
    # --------------------------------------------------------

    if show_skeleton:
        skeleton_spec = link_info.get("skeleton")
        if not isinstance(skeleton_spec, dict):
            skeleton_spec = derive_default_skeleton_from_size(size_mm)

        p1_mm = skeleton_spec.get("from", [0, 0, 0])
        p2_mm = skeleton_spec.get("to", [0, 0, 0])
        p1 = resolve_vec3_m(p1_mm, state)
        p2 = resolve_vec3_m(p2_mm, state)

        sk_radius_mm = float(skeleton_spec.get("radius", rendering_info.get("skeletonDefaults", {}).get("radius", 4)))
        sk_color = skeleton_spec.get("color", rendering_info.get("skeletonDefaults", {}).get("color", [0.85, 0.20, 0.20]))
        sk_mat = create_material_segment(f"{link_name}_skeletonMat", tuple(sk_color[:3]))

        create_cylinder_between(
            f"{link_name}_skeleton",
            p1,
            p2,
            mm_to_m(sk_radius_mm),
            link_frame,
            sk_mat
        )

    # --------------------------------------------------------
    # MARKERS
    # --------------------------------------------------------

    if show_markers:
        marker_defaults = rendering_info.get("markerDefaults", {}) or {}
        marker_mat = create_or_get_material("markerBlack")

        for m in link_info.get("markers", []):
            if not isinstance(m, dict):
                continue

            marker_name = m.get("name", f"{link_name}_marker_{m.get('id', 'x')}")
            marker_size_mm = float(m.get("size", marker_defaults.get("size", 25)))
            marker_pos = resolve_vec3_m(m.get("position", [0, 0, 0]), state)
            marker_rot = euler_deg_xyz(m.get("rotation", [0, 0, 0]))

            bpy.ops.mesh.primitive_plane_add(size=mm_to_m(marker_size_mm))
            marker_obj = bpy.context.active_object
            marker_obj.name = marker_name
            safe_parent(marker_obj, link_frame)
            marker_obj.location = marker_pos
            marker_obj.rotation_euler = marker_rot
            if len(marker_obj.data.materials) == 0:
                marker_obj.data.materials.append(marker_mat)
            else:
                marker_obj.data.materials[0] = marker_mat
# ============================================================
# 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)
)

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

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