mecha10-sim 0.8.0

//! Robot scene (.tscn) generator
//!
//! Generates standalone robot .tscn files from mecha10.json configuration.
//! The robot scene includes:
//! - CharacterBody3D root (physics body)
//! - CollisionShape3D (capsule collision)
//! - MeshInstance3D (visual representation)
//! - Sensor nodes (Camera3D, RayCast3D for LiDAR, etc.)

use anyhow::{Context, Result};
use std::fs;
use std::path::Path;

use super::profile::RobotProfile;

pub struct RobotGenerator {
    profile: RobotProfile,
    rl_mode: bool,
}

impl RobotGenerator {
    /// Create new robot generator from config file
    pub fn from_config_file(config_path: &Path) -> Result<Self> {
        let profile = RobotProfile::from_config_file(config_path)?;
        Ok(Self { profile, rl_mode: true })
    }

    /// Create robot generator with explicit RL mode setting
    pub fn from_config_file_with_mode(config_path: &Path, rl_mode: bool) -> Result<Self> {
        let profile = RobotProfile::from_config_file(config_path)?;
        Ok(Self { profile, rl_mode })
    }

    /// Generate robot .tscn file
    pub fn generate(&self, output_path: &Path) -> Result<()> {
        println!("🤖 Generating robot scene...");
        println!("   Platform: {}", self.profile.platform);
        println!("   Sensors: {}", self.profile.sensors.join(", "));

        if let Some(parent) = output_path.parent() {
            fs::create_dir_all(parent)
                .with_context(|| format!("Failed to create output directory: {}", parent.display()))?;
        }

        let tscn_content = self.generate_tscn();

        fs::write(output_path, tscn_content)
            .with_context(|| format!("Failed to write robot scene: {}", output_path.display()))?;

        println!("   ✅ Generated: {}", output_path.display());

        Ok(())
    }

    /// Generate .tscn file content
    fn generate_tscn(&self) -> String {
        let mut lines = Vec::new();

        let is_sentry = self.profile.platform == "sentry";
        let load_steps = if self.rl_mode { 5 } else { 3 } + if is_sentry { 3 } else { 0 };
        lines.push(format!("[gd_scene load_steps={} format=3]", load_steps));
        lines.push("".to_string());

        if self.rl_mode {
            lines.push(
                "[ext_resource type=\"Script\" path=\"res://components/robot_templates/rover_base_rl.gd\" id=\"1\"]"
                    .to_string(),
            );
            lines.push(
                "[ext_resource type=\"Script\" path=\"res://components/controllers/rover_ai_controller.gd\" id=\"2\"]"
                    .to_string(),
            );
            lines.push("".to_string());
        }

        lines.push("[sub_resource type=\"CapsuleShape3D\" id=\"CapsuleShape3D_1\"]".to_string());
        lines.push(format!("radius = {}", self.get_robot_radius()));
        lines.push(format!("height = {}", self.get_robot_height()));
        lines.push("".to_string());

        lines.push("[sub_resource type=\"CapsuleMesh\" id=\"CapsuleMesh_1\"]".to_string());
        lines.push(format!("radius = {}", self.get_robot_radius()));
        lines.push(format!("height = {}", self.get_robot_height()));
        lines.push("".to_string());

        if is_sentry {
            lines.push("[sub_resource type=\"BoxMesh\" id=\"PanServoMesh\"]".to_string());
            lines.push("size = Vector3(0.08, 0.06, 0.06)".to_string());
            lines.push("".to_string());

            lines.push("[sub_resource type=\"BoxMesh\" id=\"TiltServoMesh\"]".to_string());
            lines.push("size = Vector3(0.06, 0.06, 0.08)".to_string());
            lines.push("".to_string());

            lines.push("[sub_resource type=\"BoxMesh\" id=\"CameraHousingMesh\"]".to_string());
            lines.push("size = Vector3(0.05, 0.05, 0.05)".to_string());
            lines.push("".to_string());
        }

        lines.push("[node name=\"Robot\" type=\"CharacterBody3D\"]".to_string());
        if self.rl_mode {
            lines.push("script = ExtResource(\"1\")".to_string());
        }
        lines.push("".to_string());

        lines.push("[node name=\"CollisionShape3D\" type=\"CollisionShape3D\" parent=\".\"]".to_string());
        lines.push("shape = SubResource(\"CapsuleShape3D_1\")".to_string());
        lines.push("".to_string());

        lines.push("[node name=\"MeshInstance3D\" type=\"MeshInstance3D\" parent=\".\"]".to_string());
        lines.push("mesh = SubResource(\"CapsuleMesh_1\")".to_string());
        lines.push("".to_string());

        if self.rl_mode {
            lines.push("[node name=\"AIController3D\" type=\"Node3D\" parent=\".\"]".to_string());
            lines.push("script = ExtResource(\"2\")".to_string());
            lines.push("".to_string());
        }

        self.add_sensor_nodes(&mut lines);

        lines.push("[node name=\"Markers\" type=\"Node3D\" parent=\".\"]".to_string());
        lines.push("".to_string());

        lines.push("[node name=\"CenterOfMass\" type=\"Marker3D\" parent=\"Markers\"]".to_string());
        lines.push("".to_string());

        lines.join("\n")
    }

    /// Add sensor nodes based on robot profile
    fn add_sensor_nodes(&self, lines: &mut Vec<String>) {
        let parent = ".";

        let is_sentry = self.profile.platform == "sentry";

        if is_sentry {
            // Sentry has no flat, fixed-angle camera mount - it mounts its camera on an
            // articulated pan/tilt joint chain instead. See `add_pan_tilt_rig`.
            self.add_pan_tilt_rig(lines);
        } else if self.profile.sensors.contains(&"camera".to_string()) {
            if self.rl_mode {
                lines.push(format!(
                    "[node name=\"CameraRL\" type=\"Node3D\" parent=\"{}\"]",
                    parent
                ));
                lines.push("script = ExtResource(\"res://components/sensors/camera_rl/camera_rl.gd\")".to_string());
                lines.push(format!(
                    "transform = Transform3D(1, 0, 0, 0, 1, 0, 0, 0, 1, 0, {}, {})",
                    self.get_robot_height() / 2.0 - 0.1,
                    0.3
                ));
                lines.push("resolution = Vector2i(84, 84)".to_string());
                lines.push("fov_degrees = 90.0".to_string());
                lines.push("observation_enabled = true".to_string());
                lines.push("".to_string());
            } else {
                lines.push(format!(
                    "[node name=\"Camera\" type=\"Camera3D\" parent=\"{}\"]",
                    parent
                ));
                lines.push(format!(
                    "transform = Transform3D(1, 0, 0, 0, 1, 0, 0, 0, 1, 0, {}, {})",
                    self.get_robot_height() / 2.0 - 0.1,
                    0.3
                ));
                lines.push("fov = 60.0".to_string());
                lines.push("".to_string());
            }
        }

        if self.profile.sensors.contains(&"lidar".to_string()) {
            if self.rl_mode {
                lines.push(format!("[node name=\"LidarRL\" type=\"Node3D\" parent=\"{}\"]", parent));
                lines.push("script = ExtResource(\"res://components/sensors/lidar_rl/lidar_rl.gd\")".to_string());
                lines.push(format!(
                    "transform = Transform3D(1, 0, 0, 0, 1, 0, 0, 0, 1, 0, {}, 0)",
                    self.get_robot_height() / 2.0
                ));
                lines.push("num_rays = 16".to_string());
                lines.push("ray_length = 10.0".to_string());
                lines.push("observation_enabled = true".to_string());
                lines.push("".to_string());
            } else {
                lines.push(format!("[node name=\"LiDAR\" type=\"Node3D\" parent=\"{}\"]", parent));
                lines.push(format!(
                    "transform = Transform3D(1, 0, 0, 0, 1, 0, 0, 0, 1, 0, {}, 0)",
                    self.get_robot_height() / 2.0
                ));
                lines.push("".to_string());

                for i in 0..36 {
                    let angle_deg = i * 10;
                    let angle_rad = (angle_deg as f32).to_radians();
                    let x = angle_rad.cos();
                    let z = angle_rad.sin();

                    lines.push(format!(
                        "[node name=\"Ray{}\" type=\"RayCast3D\" parent=\"{}/LiDAR\"]",
                        i, parent
                    ));
                    lines.push(format!("target_position = Vector3({}, 0, {})", x * 10.0, z * 10.0));
                    lines.push("enabled = true".to_string());
                    lines.push("".to_string());
                }
            }
        }

        if self.profile.sensors.contains(&"imu".to_string()) {
            lines.push(format!("[node name=\"IMU\" type=\"Marker3D\" parent=\"{}\"]", parent));
            lines.push("".to_string());
        }
    }

    /// Add the sentry platform's pan (yaw) / tilt (pitch) joint chain, with its camera mounted
    /// at the end of the chain instead of as a flat, fixed-angle child of the robot root.
    ///
    /// Sentry is a stationary base (no locomotion) whose camera is carried by two nested
    /// pivots: `PanPivot` (parented at the robot root, sitting on top of the base) and
    /// `TiltPivot` (parented at `PanPivot`, so tilt rotates relative to whatever pan has
    /// already rotated to). Each pivot gets its own `BoxMesh` servo body so the rig is
    /// actually visible instead of relying on editor-only gizmos for empty `Node3D`s, and
    /// the tilt servo sits a bit above the pan servo so it has room to rotate without
    /// clipping into it. The camera gets a housing mesh as a sibling under `TiltPivot` so
    /// it visibly reads as mounted on the top (tilt) servo rather than floating detached.
    /// The camera itself is only emitted when the profile actually lists a `"camera"`
    /// sensor - the pivots and their servo meshes are emitted unconditionally since they
    /// represent the robot's real joints regardless of what sensors are configured.
    fn add_pan_tilt_rig(&self, lines: &mut Vec<String>) {
        lines.push("[node name=\"PanPivot\" type=\"Node3D\" parent=\".\"]".to_string());
        lines.push(format!(
            "transform = Transform3D(1, 0, 0, 0, 1, 0, 0, 0, 1, 0, {}, 0)",
            self.get_robot_height() / 2.0
        ));
        lines.push("".to_string());

        lines.push("[node name=\"PanServoMesh\" type=\"MeshInstance3D\" parent=\"PanPivot\"]".to_string());
        lines.push("mesh = SubResource(\"PanServoMesh\")".to_string());
        lines.push("".to_string());

        lines.push("[node name=\"TiltPivot\" type=\"Node3D\" parent=\"PanPivot\"]".to_string());
        lines.push("".to_string());

        lines.push("[node name=\"TiltServoMesh\" type=\"MeshInstance3D\" parent=\"PanPivot/TiltPivot\"]".to_string());
        lines.push("mesh = SubResource(\"TiltServoMesh\")".to_string());
        lines.push("transform = Transform3D(1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0.05, 0)".to_string());
        lines.push("".to_string());

        lines.push(
            "[node name=\"CameraHousingMesh\" type=\"MeshInstance3D\" parent=\"PanPivot/TiltPivot\"]".to_string(),
        );
        lines.push("mesh = SubResource(\"CameraHousingMesh\")".to_string());
        lines.push("transform = Transform3D(1, 0, 0, 0, 1, 0, 0, 0, 1, 0.08, 0.05, 0)".to_string());
        lines.push("".to_string());

        if self.profile.sensors.contains(&"camera".to_string()) {
            if self.rl_mode {
                lines.push("[node name=\"CameraRL\" type=\"Node3D\" parent=\"PanPivot/TiltPivot\"]".to_string());
                lines.push("script = ExtResource(\"res://components/sensors/camera_rl/camera_rl.gd\")".to_string());
                lines.push("transform = Transform3D(1, 0, 0, 0, 1, 0, 0, 0, 1, 0.08, 0.05, 0)".to_string());
                lines.push("resolution = Vector2i(84, 84)".to_string());
                lines.push("fov_degrees = 90.0".to_string());
                lines.push("observation_enabled = true".to_string());
                lines.push("".to_string());
            } else {
                lines.push("[node name=\"Camera\" type=\"Camera3D\" parent=\"PanPivot/TiltPivot\"]".to_string());
                lines.push("transform = Transform3D(1, 0, 0, 0, 1, 0, 0, 0, 1, 0.08, 0.05, 0)".to_string());
                lines.push("fov = 60.0".to_string());
                lines.push("".to_string());
            }
        }
    }

    /// Per-platform (radius, height) dimensions, in meters. Shared table behind
    /// [`Self::get_robot_radius`]/[`Self::get_robot_height`], which only differ in which
    /// element of the pair they return.
    fn platform_dimensions(platform: &str) -> (f32, f32) {
        match platform {
            "differential-drive" => (0.3, 1.0),
            "holonomic" => (0.35, 1.2),
            "legged" => (0.25, 1.5),
            "arm" => (0.15, 2.0),
            "sentry" => (0.1, 0.4),
            _ => (0.3, 1.0),
        }
    }

    /// Get robot radius based on platform
    fn get_robot_radius(&self) -> f32 {
        Self::platform_dimensions(self.profile.platform.as_str()).0
    }

    /// Get robot height based on platform
    fn get_robot_height(&self) -> f32 {
        Self::platform_dimensions(self.profile.platform.as_str()).1
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_profile(platform: &str, sensors: &[&str]) -> RobotProfile {
        RobotProfile {
            platform: platform.to_string(),
            sensors: sensors.iter().map(|s| s.to_string()).collect(),
            sensor_configs: Vec::new(),
            task_nodes: Vec::new(),
            capabilities: Vec::new(),
            difficulty_preference: "medium".to_string(),
        }
    }

    #[test]
    fn sentry_platform_dimensions_are_compact_not_differential_drive_fallback() {
        let sentry_dims = RobotGenerator::platform_dimensions("sentry");
        let fallback_dims = RobotGenerator::platform_dimensions("differential-drive");

        assert_eq!(sentry_dims, (0.1, 0.4));
        assert_ne!(sentry_dims, fallback_dims);
    }

    #[test]
    fn sentry_with_camera_emits_pan_tilt_rig_and_camera() {
        let profile = make_profile("sentry", &["camera"]);
        let generator = RobotGenerator {
            profile,
            rl_mode: false,
        };

        let tscn = generator.generate_tscn();

        assert!(tscn.contains("PanPivot"));
        assert!(tscn.contains("TiltPivot"));
        assert!(tscn.contains("Camera"));
        assert!(tscn.contains("parent=\"PanPivot/TiltPivot\""));
    }

    #[test]
    fn non_sentry_platform_never_emits_pan_tilt_rig() {
        let profile = make_profile("differential-drive", &["camera"]);
        let generator = RobotGenerator {
            profile,
            rl_mode: false,
        };

        let tscn = generator.generate_tscn();

        assert!(!tscn.contains("PanPivot"));
        assert!(!tscn.contains("TiltPivot"));
    }

    #[test]
    fn sentry_without_camera_still_emits_rig_but_no_camera_node() {
        let profile = make_profile("sentry", &["lidar"]);
        let generator = RobotGenerator {
            profile,
            rl_mode: false,
        };

        let tscn = generator.generate_tscn();

        assert!(tscn.contains("PanPivot"));
        assert!(tscn.contains("TiltPivot"));
        // The pan/tilt rig's `CameraHousingMesh` is emitted unconditionally (it's cosmetic
        // mount hardware, not a sensor), so it legitimately contains the substring "Camera".
        // What must be absent is the actual `Camera3D` sensor node, since no "camera" sensor
        // was configured on this profile.
        assert!(!tscn.contains("type=\"Camera3D\""));
    }
}