import json
import os
import tempfile
import xml.etree.ElementTree as ET

from ament_index_python.packages import get_package_share_directory
from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument, IncludeLaunchDescription, OpaqueFunction, SetEnvironmentVariable
from launch.conditions import IfCondition
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch.substitutions import LaunchConfiguration
from launch_ros.actions import Node


HOUSE_FURNITURE_BOXES = {
    # mailbox.dae uses inch units and an internal transform. These values preserve
    # its rendered bounds while keeping the visual box upright on the floor.
    'mailbox/': ('0.133 0.355 1.067', '0.002 0.017 0.533 0 0 0'),
    'cafe_table/': ('0.94 0.94 0.78', None),
    'first_2015_trash_can/': ('0.44 0.44 0.76', None),
    'table_marble/': ('1.40 0.86 0.76', None),
}

LOCAL_GROUND_PLANE_MODEL = '''    <model name="ground_plane">
      <static>1</static>
      <link name="link">
        <collision name="collision"><geometry><plane><normal>0 0 1</normal><size>100 100</size></plane></geometry></collision>
        <visual name="visual"><cast_shadows>false</cast_shadows><geometry><plane><normal>0 0 1</normal><size>100 100</size></plane></geometry><material><ambient>0.24 0.27 0.32 1</ambient><diffuse>0.24 0.27 0.32 1</diffuse></material></visual>
      </link>
    </model>'''

LOCAL_SUN_LIGHT = '''    <light type="directional" name="sun">
      <cast_shadows>false</cast_shadows><pose>0 0 10 0 0 0</pose><diffuse>0.8 0.8 0.8 1</diffuse><specular>0.2 0.2 0.2 1</specular><direction>-0.5 0.1 -0.9</direction>
    </light>'''

MKMINI_WORLD_SPACING_SCALE = 1.0


def replace_house_visual_meshes_with_boxes(house_model):
    """Keep physical mesh collisions, but render furniture as inexpensive boxes."""
    for visual in house_model.findall('.//visual'):
        geometry = visual.find('geometry')
        mesh = geometry.find('mesh') if geometry is not None else None
        uri = mesh.findtext('uri', default='') if mesh is not None else ''
        box_config = next((value for key, value in HOUSE_FURNITURE_BOXES.items() if key in uri), None)
        if box_config is None:
            continue

        size, pose = box_config
        geometry.remove(mesh)
        box = ET.SubElement(geometry, 'box')
        ET.SubElement(box, 'size').text = size
        if pose is not None:
            visual_pose = visual.find('pose')
            if visual_pose is None:
                visual_pose = ET.SubElement(visual, 'pose')
            visual_pose.text = pose

    # The house references Gazebo's external mailbox asset, which is not shipped
    # with the module image. Replace both its visual and collision with the same
    # local box so it remains visible and navigable in the flat indoor world.
    mailbox = next(
        (model for model in house_model.findall('.//model') if model.attrib.get('name') == 'Mailbox'),
        None,
    )
    if mailbox is None:
        return

    mailbox_link = mailbox.find('link')
    if mailbox_link is None:
        return

    for element in list(mailbox_link):
        if element.tag in {'visual', 'collision'}:
            mailbox_link.remove(element)

    # The laser is roughly 17 cm above the ground. Keep the mailbox body above
    # that plane and retain the single thin support from mailbox.dae.
    mailbox_parts = (
        ('body', '0.133 0.355 0.29', '0.002 0.017 0.922 0 0 0'),
        ('leg', '0.043 0.044 0.803', '0.000 0.012 0.401 0 0 0'),
    )
    for element_type in ('collision', 'visual'):
        for name, size, pose in mailbox_parts:
            element = ET.SubElement(mailbox_link, element_type, name=f'mailbox_{element_type}_{name}')
            ET.SubElement(element, 'pose').text = pose
            geometry = ET.SubElement(element, 'geometry')
            box = ET.SubElement(geometry, 'box')
            ET.SubElement(box, 'size').text = size


def create_flat_house_world(gazebo_dir):
    """Create a texture-free house copy with box-rendered furniture."""
    world_path = os.path.join(gazebo_dir, 'worlds', 'turtlebot3_house.world')
    model_path = os.path.join(gazebo_dir, 'models', 'turtlebot3_house', 'model.sdf')
    flat_models_dir = os.path.join(tempfile.gettempdir(), 'nav2_flat_models')
    flat_model_dir = os.path.join(flat_models_dir, 'turtlebot3_house_flat')
    os.makedirs(flat_model_dir, exist_ok=True)

    house_model_tree = ET.parse(model_path)
    house_sdf = house_model_tree.getroot()
    house_model = house_sdf.find('model')
    if house_model is None:
        raise RuntimeError('turtlebot3_house model is missing from model.sdf')
    house_model.attrib['name'] = 'turtlebot3_house_flat'
    house_static = house_model.find('static')
    if house_static is None:
        house_static = ET.Element('static')
        house_model.insert(0, house_static)
    house_static.text = '1'
    replace_house_visual_meshes_with_boxes(house_model)

    for visual in house_model.findall('.//visual'):
        material = visual.find('material')
        if material is None:
            material = ET.SubElement(visual, 'material')

        for script in material.findall('script'):
            material.remove(script)

        for name, value in (
            ('ambient', '0.62 0.67 0.76 1'),
            ('diffuse', '0.62 0.67 0.76 1'),
            ('specular', '0.05 0.05 0.05 1'),
        ):
            color = material.find(name)
            if color is None:
                color = ET.SubElement(material, name)
            color.text = value

        shadows = visual.find('cast_shadows')
        if shadows is None:
            shadows = ET.SubElement(visual, 'cast_shadows')
        shadows.text = 'false'

        lighting = visual.find('lighting')
        if lighting is None:
            lighting = ET.SubElement(visual, 'lighting')
        lighting.text = 'false'

    # ElementTree serializes Gazebo's frame='' as frame="", which is invalid SDF.
    for element in house_sdf.iter():
        if element.attrib.get('frame') == '':
            del element.attrib['frame']

    house_model_tree.write(
        os.path.join(flat_model_dir, 'model.sdf'),
        encoding='utf-8',
        xml_declaration=True,
    )
    with open(os.path.join(flat_model_dir, 'model.config'), 'w', encoding='utf-8') as config:
        config.write(
            '<model><name>turtlebot3_house_flat</name>'
            '<version>1.0</version><sdf version="1.6">model.sdf</sdf></model>'
        )

    with open(world_path, 'r', encoding='utf-8') as source:
        flat_world = source.read()
    flat_world = flat_world.replace(
        '    <include>\n      <uri>model://ground_plane</uri>\n    </include>',
        LOCAL_GROUND_PLANE_MODEL,
    ).replace(
        '    <include>\n      <uri>model://sun</uri>\n    </include>',
        LOCAL_SUN_LIGHT,
    ).replace('model://turtlebot3_house', 'model://turtlebot3_house_flat')

    flat_world_path = os.path.join(tempfile.gettempdir(), 'nav2_turtlebot3_house_flat.world')
    with open(flat_world_path, 'w', encoding='utf-8') as destination:
        destination.write(flat_world)

    existing_model_path = os.environ.get('GAZEBO_MODEL_PATH', '')
    os.environ['GAZEBO_MODEL_PATH'] = os.pathsep.join(
        part for part in (flat_models_dir, existing_model_path) if part
    )
    os.environ['GAZEBO_MODEL_DATABASE_URI'] = ''
    return flat_world_path


def create_mkmini_world(gazebo_dir):
    """Create the MKMini world variant without changing map-scale coordinates."""
    source_model_path = os.path.join(gazebo_dir, 'models', 'turtlebot3_world', 'model.sdf')
    models_dir = os.path.join(tempfile.gettempdir(), 'nav2_mkmini_models')
    model_dir = os.path.join(models_dir, 'turtlebot3_world_mkmini')
    os.makedirs(model_dir, exist_ok=True)

    model_tree = ET.parse(source_model_path)
    model_sdf = model_tree.getroot()
    model = model_sdf.find('model')
    if model is None:
        raise RuntimeError('turtlebot3_world model is missing from model.sdf')
    model.attrib['name'] = 'turtlebot3_world_mkmini'

    # Keep object geometry and coordinates aligned with the base world. MKMini
    # clearance comes from its footprint and inflation configuration.
    for pose in model.findall('.//pose'):
        values = pose.text.split()
        if len(values) != 6:
            continue
        try:
            values[0] = str(float(values[0]) * MKMINI_WORLD_SPACING_SCALE)
            values[1] = str(float(values[1]) * MKMINI_WORLD_SPACING_SCALE)
        except ValueError:
            continue
        pose.text = ' '.join(values)

    # The black wall is centered at the origin, so its pose is unaffected by
    # the spacing expansion. Enlarge both its visual and collision mesh to
    # retain the same relationship with the widened green obstacles.
    for mesh in model.findall('.//mesh'):
        if not (mesh.findtext('uri', default='').endswith('/wall.dae')):
            continue
        scale = mesh.find('scale')
        if scale is None:
            continue
        values = scale.text.split()
        if len(values) != 3:
            continue
        try:
            scale.text = ' '.join(str(float(value) * MKMINI_WORLD_SPACING_SCALE) for value in values)
        except ValueError:
            continue

    # Keep only the centre post from TurtleBot3 World's 3x3 white-pillar grid.
    # This is applied to the generated MKMini world only; Burger keeps the
    # original obstacle layout. Remove both visual and collision shapes so the
    # Gazebo scene and Nav2 costmap describe the same space.
    for link in model.findall('.//link'):
        for element in list(link):
            if element.tag not in {'collision', 'visual'}:
                continue
            cylinder = element.find('./geometry/cylinder')
            if cylinder is None:
                continue
            if cylinder.findtext('radius') != '0.15' or cylinder.findtext('length') != '0.5':
                continue
            pose = element.find('pose')
            values = pose.text.split() if pose is not None and pose.text else []
            if len(values) != 6:
                continue
            try:
                x, y = float(values[0]), float(values[1])
            except ValueError:
                continue
            if abs(x) > 1e-6 or abs(y) > 1e-6:
                link.remove(element)
                continue
            # The MKMini lidar is centered at 0.52 m. Raise the remaining
            # centre post so it stays observable instead of below the scan.
            cylinder.find('length').text = '1.2'
            values[2] = '0.6'
            pose.text = ' '.join(values)

    model_tree.write(
        os.path.join(model_dir, 'model.sdf'),
        encoding='utf-8',
        xml_declaration=True,
    )
    with open(os.path.join(model_dir, 'model.config'), 'w', encoding='utf-8') as config:
        config.write(
            '<model><name>turtlebot3_world_mkmini</name>'
            '<version>1.0</version><sdf version="1.6">model.sdf</sdf></model>'
        )

    source_world_path = os.path.join(gazebo_dir, 'worlds', 'turtlebot3_world.world')
    with open(source_world_path, 'r', encoding='utf-8') as source:
        world = source.read().replace('model://turtlebot3_world', 'model://turtlebot3_world_mkmini')
    world_path = os.path.join(tempfile.gettempdir(), 'nav2_mkmini_world.world')
    with open(world_path, 'w', encoding='utf-8') as destination:
        destination.write(world)
    return world_path, models_dir


def create_custom_layout_world():
    """Build a Gazebo world from the editor's persisted 2D layout."""
    layout_path = os.environ.get('WORLD_LAYOUT_PATH', '/modules/world-layout.json')
    try:
        with open(layout_path, 'r', encoding='utf-8') as source:
            objects = json.load(source).get('objects', [])
    except (OSError, ValueError, AttributeError):
        objects = []

    models = []
    for index, item in enumerate(objects):
        if not isinstance(item, dict):
            continue
        position = item.get('position') or [0, 0, 0]
        rotation = item.get('rotation') or [0, 0, 0]
        size = item.get('size') or [1, 1, 1]
        try:
            x, y, z = (float(position[0]), float(position[1]), float(position[2]))
            yaw = float(rotation[2])
            sx, sy, sz = (max(0.05, float(size[0])), max(0.05, float(size[1])), max(0.05, float(size[2])))
        except (IndexError, TypeError, ValueError):
            continue

        name = f"custom_{item.get('type', 'box')}_{index}"
        if item.get('type') == 'pillar':
            geometry = f'<cylinder><radius>{sx / 2.0}</radius><length>{sz}</length></cylinder>'
        else:
            geometry = f'<box><size>{sx} {sy} {sz}</size></box>'
        models.append(
            f'''<model name="{name}">
  <static>true</static>
  <pose>{x} {y} {z} 0 0 {yaw}</pose>
  <link name="link">
    <collision name="collision"><geometry>{geometry}</geometry></collision>
    <visual name="visual"><geometry>{geometry}</geometry><material><ambient>0.34 0.58 0.82 1</ambient><diffuse>0.34 0.58 0.82 1</diffuse></material></visual>
  </link>
</model>'''
        )

    world_path = os.path.join(tempfile.gettempdir(), 'nav2_custom_layout.world')
    with open(world_path, 'w', encoding='utf-8') as destination:
        destination.write('''<?xml version="1.0"?>
<sdf version="1.6">
  <world name="nav2_custom_layout">
    <include><uri>model://sun</uri></include>
    <include><uri>model://ground_plane</uri></include>
    <physics type="ode"><real_time_update_rate>1000</real_time_update_rate></physics>
''')
        destination.write('\n'.join(models))
        destination.write('\n  </world>\n</sdf>\n')
    return world_path


def launch_gazebo(context):
    gazebo_dir = get_package_share_directory('turtlebot3_gazebo')
    simulator_dir = get_package_share_directory('nav2_turtlebot_simulator')
    gazebo_ros_dir = get_package_share_directory('gazebo_ros')
    world = LaunchConfiguration('world').perform(context)
    render_mode = LaunchConfiguration('render_mode').perform(context).strip().lower()
    platform = LaunchConfiguration('platform').perform(context).strip().lower()
    spawn_x = LaunchConfiguration('x_pose').perform(context)
    spawn_y = LaunchConfiguration('y_pose').perform(context)
    extra_model_path = ''

    if render_mode not in {'flat', 'full'}:
        raise RuntimeError('render_mode must be either flat or full')

    if world == 'empty_world.world' and os.path.isfile(os.environ.get('WORLD_LAYOUT_PATH', '/modules/world-layout.json')):
        world_path = create_custom_layout_world()
    elif world == 'turtlebot3_house.world' and render_mode == 'flat':
        world_path = create_flat_house_world(gazebo_dir)
    elif world == 'turtlebot3_world.world' and platform in {'mkmini', 'mk-mini'}:
        world_path, extra_model_path = create_mkmini_world(gazebo_dir)
    else:
        world_path = os.path.join(
            simulator_dir if world == 'empty_world.world' else gazebo_dir,
            'worlds',
            world,
        )

    if platform in {'turtlebot3', 'turtlebot3_burger', 'burger'}:
        entity_name = 'turtlebot3_burger_3d_lidar'
        model_path = os.path.join(gazebo_dir, 'models', 'turtlebot3_burger', 'model.sdf')
        model_directory = os.path.join(gazebo_dir, 'models')
        static_transforms = [
            Node(
                package='tf2_ros',
                executable='static_transform_publisher',
                name='turtlebot3_base_link_tf',
                arguments=[
                    '--x', '0', '--y', '0', '--z', '0.01',
                    '--roll', '0', '--pitch', '0', '--yaw', '0',
                    '--frame-id', 'base_footprint', '--child-frame-id', 'base_link',
                ],
            ),
            Node(
                package='tf2_ros',
                executable='static_transform_publisher',
                name='turtlebot3_laser_tf',
                arguments=[
                    '--x', '-0.032', '--y', '0', '--z', '0.171',
                    '--roll', '0', '--pitch', '0', '--yaw', '0',
                    '--frame-id', 'base_link', '--child-frame-id', 'base_scan',
                ],
            ),
        ]
    elif platform in {'mkmini', 'mk-mini'}:
        entity_name = 'mkmini_gazebo'
        model_directory = os.path.join(simulator_dir, 'models')
        model_path = os.path.join(model_directory, 'mkmini_gazebo', 'model.sdf')
        # The world obstacles are spread out for MKMini, but the original
        # TurtleBot spawn is already inside the navigable centre of that map.
        # Scaling it as well placed the robot against the generated map edge.
        if spawn_x == '-2.0' and spawn_y == '-0.5':
            spawn_x = '-1.4'
        # Gazebo's fixed joint does not reliably publish this transform to the
        # separate SLAM container. Publish the sensor mounting point explicitly.
        static_transforms = [
            Node(
                package='tf2_ros',
                executable='static_transform_publisher',
                name='mkmini_lidar_tf',
                arguments=[
                    '--x', '0.15', '--y', '0', '--z', '0.578',
                    '--roll', '0', '--pitch', '0', '--yaw', '0',
                    '--frame-id', 'base_link', '--child-frame-id', 'os_lidar',
                ],
            ),
        ]
    else:
        raise RuntimeError('platform must be turtlebot3 or mkmini')

    existing_model_path = os.environ.get('GAZEBO_MODEL_PATH', '')
    return [
        SetEnvironmentVariable(
            name='GAZEBO_MODEL_PATH',
            value=os.pathsep.join(part for part in (extra_model_path, model_directory, existing_model_path) if part),
        ),
        IncludeLaunchDescription(
            PythonLaunchDescriptionSource(os.path.join(gazebo_ros_dir, 'launch', 'gzserver.launch.py')),
            launch_arguments={
                'world': world_path,
                'extra_gazebo_args': '-slibgazebo_ros_state.so',
            }.items(),
        ),
        Node(
            package='gazebo_ros',
            executable='spawn_entity.py',
            arguments=[
                '-entity', entity_name,
                '-file', model_path,
                '-x', spawn_x,
                '-y', spawn_y,
                '-z', '0.01',
            ],
            output='screen',
        ),
        *static_transforms,
    ]


def generate_launch_description():
    gazebo_ros_dir = get_package_share_directory('gazebo_ros')
    return LaunchDescription([
        DeclareLaunchArgument('gui', default_value='false'),
        DeclareLaunchArgument('world', default_value='turtlebot3_world.world'),
        DeclareLaunchArgument('render_mode', default_value='flat'),
        DeclareLaunchArgument('platform', default_value='turtlebot3'),
        DeclareLaunchArgument('x_pose', default_value='-2.0'),
        DeclareLaunchArgument('y_pose', default_value='-0.5'),
        OpaqueFunction(function=launch_gazebo),
        IncludeLaunchDescription(
            PythonLaunchDescriptionSource(os.path.join(gazebo_ros_dir, 'launch', 'gzclient.launch.py')),
            condition=IfCondition(LaunchConfiguration('gui')),
        ),
    ])