import json
import re
from pathlib import Path
from typing import Dict, Any, List
from ..base import BaseBridge

class IsaacBridge(BaseBridge):
    """
    Transforms OpenRGD definitions into Isaac Lab configuration classes (Python).
    Generates 'robot_cfg.py' ready for Reinforcement Learning environments.
    """
    
    def generate(self, output_dir: Path) -> None:
        self.log(f"Transpiling identity {self.robot_id} to Isaac Lab format...")
        
        # 1. LOAD TWIN
        unified_path = self.spec_dir / "openrgd_unified_spec.json"
        if not unified_path.exists():
            self.log("❌ Machine Twin not found. Run 'rgd compile-spec' first.")
            return

        try:
            with open(unified_path, 'r', encoding='utf-8') as f:
                unified_data = json.load(f)
        except Exception as e:
            self.log(f"❌ Error: {e}")
            return

        # 2. EXTRACT MODULES
        actuation_topology = None
        description = None
        
        for module in unified_data.get("files", []):
            mid = module.get("id")
            if mid == "actuation_topology": actuation_topology = module.get("content")
            elif mid == "description": description = module.get("content")

        if not actuation_topology:
            self.log("❌ Critical: 'actuation_topology' missing.")
            return

        # 3. GROUP JOINTS BY PROFILE (The Clean Code Strategy)
        # We want to generate: "HEAVY_LEGS = ImplicitActuatorCfg(...)"
        # Instead of 20 separate configs.
        
        profiles_map = actuation_topology.get("control_profiles_map", {})
        instances_map = actuation_topology.get("joint_actuator_mapping_map", {})
        
        # Grouping structure: { 'profile_name': { 'joints': [], 'config': {...} } }
        groups = {}
        
        for joint_name, data in instances_map.items():
            profile_ref = data.get("use_profile_ref_str", "custom_unnamed")
            
            # If it's a custom group or override, we treat it separately
            if "overrides" in data:
                profile_ref = f"custom_{joint_name}"
                # Merge logic would go here, simplified for v0.1
                config = data.get("control_defaults", {})
            else:
                # Use profile data
                config = profiles_map.get(profile_ref, {}).get("control_defaults", {})
            
            if profile_ref not in groups:
                groups[profile_ref] = {"joints": [], "config": config}
            
            groups[profile_ref]["joints"].append(joint_name)

        # 4. GENERATE PYTHON CODE
        self._write_python_config(groups, description, output_dir)

    def _write_python_config(self, groups: Dict, description: Dict, output_dir: Path):
        lines = []
        
        # Imports
        lines.append(f"# ISAAC LAB CONFIGURATION - GENERATED BY OPENRGD")
        lines.append(f"# Robot ID: {self.robot_id}")
        lines.append("")
        lines.append("from isaaclab.actuators import ImplicitActuatorCfg")
        lines.append("from isaaclab.assets import ArticulationCfg")
        lines.append("import isaaclab.sim as sim_utils")
        lines.append("")
        lines.append(f"class {self._to_pascal(self.robot_id)}Cfg(ArticulationCfg):")
        lines.append("    \"\"\"Configuration for the robot.\"\"\"")
        lines.append("")
        
        # Spawn Config (USD Path)
        usd_path = "assets/robot.usd" # Fallback
        if description:
            # Try to find USD path in description
            # Note: deeply nested check omitted for brevity
            pass
            
        lines.append("    spawn = sim_utils.UsdFileCfg(")
        lines.append(f"        usd_path=\"{usd_path}\",")
        lines.append("        activate_contact_sensors=True,")
        lines.append("        rigid_props=sim_utils.RigidBodyPropertiesCfg(")
        lines.append("            disable_gravity=False,")
        lines.append("        ),")
        lines.append("    )")
        lines.append("")
        lines.append("    init_state = ArticulationCfg.InitialStateCfg(")
        lines.append("        pos=(0.0, 0.0, 0.0)")
        lines.append("    )")
        lines.append("")
        lines.append("    # --- ACTUATORS ---")
        lines.append("    # Grouped by OpenRGD Control Profiles")
        lines.append("")

        # Generate Actuator Groups
        for profile_name, group_data in groups.items():
            config = group_data["config"]
            joint_names = group_data["joints"]
            
            # Extract Isaac Params
            imp = config.get("advanced_impedance_model", {})
            stiffness = imp.get("stiffness_nm_rad_float") or imp.get("stiffness", 0.0)
            damping = imp.get("damping_nms_rad_float") or imp.get("damping", 0.0)
            
            # Convert OpenRGD Regex to Isaac Regex
            # Isaac uses regex list. We just list explicit names for safety here.
            # To make it cleaner, we join them with OR '|' if there are many
            regex_pattern = "|".join(joint_names)
            
            sanitized_name = profile_name.upper().replace("PROFILE_", "")
            
            lines.append(f"    {sanitized_name} = ImplicitActuatorCfg(")
            lines.append(f"        joint_names_expr=[\"{regex_pattern}\"],")
            lines.append(f"        stiffness={stiffness},")
            lines.append(f"        damping={damping},")
            lines.append(f"        effort_limit=None,  # Loaded from USD")
            lines.append(f"        velocity_limit=None, # Loaded from USD")
            lines.append(f"    )")
            lines.append("")

        # Write File
        out_file = output_dir / "isaac_robot_cfg.py"
        with open(out_file, "w", encoding="utf-8") as f:
            f.write("\n".join(lines))
            
        self.log(f"✅ Generated: {out_file.name}")

    def _to_pascal(self, snake_str):
        # Helper: "berkeley-humanoid" -> "BerkeleyHumanoid"
        components = snake_str.replace("-", "_").split('_')
        return ''.join(x.title() for x in components)