import os
from pathlib import Path

from ament_index_python.packages import get_package_share_directory

from launch.substitutions import Command

from launch.actions import IncludeLaunchDescription, TimerAction
from launch import LaunchDescription
from launch.actions import IncludeLaunchDescription
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch.substitutions import LaunchConfiguration
from launch.actions import DeclareLaunchArgument


from launch_ros.actions import Node
from launch.actions import RegisterEventHandler
from launch.event_handlers import OnProcessExit

from pathlib import Path
from typing import List, Tuple, Union, cast, Dict
from struct import Struct
from enum import IntEnum

import numpy as np



class ConeTypes(IntEnum):
    """
    Enum for all possible cone types
    """

    UNKNOWN = 0
    YELLOW = 1
    RIGHT = 1
    BLUE = 2
    LEFT = 2
    ORANGE_SMALL = 3
    START_FINISH_AREA = 3
    ORANGE_BIG = 4
    START_FINISH_LINE = 4


HEADER_STRUCT = Struct("6sBBhBB")
BLOCK_STRUCT = Struct("2h4B")

LytObjectIndexToConeType: Dict[int, ConeTypes] = {
    25: ConeTypes.UNKNOWN,
    29: ConeTypes.YELLOW,
    30: ConeTypes.YELLOW,
    23: ConeTypes.BLUE,
    24: ConeTypes.BLUE,
    27: ConeTypes.ORANGE_BIG,
    20: ConeTypes.ORANGE_SMALL,
}   

def split_header_blocks(data: bytes) -> Tuple[bytes, bytes]:
    """
    Split the content of the lyt file into header and block. This split is easy because
    the header has a fixed size

    Args:
        data (bytes): The content of the lyt file

    Returns:
        Tuple[bytes, bytes]: The header and the block
    """
    return data[: HEADER_STRUCT.size], data[HEADER_STRUCT.size :]


def verify_lyt_header(header_data: bytes) -> None:
    """
    Parse the header and perform some sanity checks suggested by the LFS documentation

    Args:
        header_data (bytes): The header bytes of the `.lyt` file
    """

    header = cast(
        Tuple[bytes, int, int, int, int, int], HEADER_STRUCT.unpack(header_data)
    )

    file_type, version, revision, _, _, _ = header
    assert file_type == b"LFSLYT"
    assert version <= 0, version
    # revision allowed up to 252
    # https://www.lfs.net/forum/thread/96153-LYT-revision-252-in-W-patch
    assert revision <= 252, revision


def extract_cone_lists(blocks_data: bytes) -> List[List[Tuple[float, float]]]:
    """
    Extract the cone object positions from the object blocks bytes of a lyt file

    Args:
        blocks_data (bytes): The data in the lyt file that is not the header

    Returns:
        List[List[Tuple[int, int]]]: The cone positions split by cone type
    """
    decoded_blocks = BLOCK_STRUCT.iter_unpack(blocks_data)
    all_cones_per_type: List[List[Tuple[float, float]]] = [[] for _ in ConeTypes]

    # cone_info:
    for cone_info in decoded_blocks:
        obj_x, obj_y, _, _, lyt_obj_idx, _ = cast(
            Tuple[int, int, int, int, int, int], cone_info
        )

        try:
            cone_type = LytObjectIndexToConeType[lyt_obj_idx]
        except KeyError:
            # not a cone
            continue

        # the stored x,y pos is multiplied by
        # 16 in the file so we need to convert it back
        # (and cast to a float by using real div)
        obj_x_meters = obj_x / 16
        obj_y_meters = obj_y / 16
        all_cones_per_type[cone_type].append((obj_x_meters, obj_y_meters))
    return all_cones_per_type


def load_lyt_file(filename: Union[Path, str]) -> List[np.ndarray]:
    """
    Load a `.lyt` file and return the positions of the cone objects inside it split
    according to `ConeTypes`

    Args:
        filename (Path): The path to the `.lyt` file

    Returns:
        List[np.ndarray]: A list of 2d np.ndarrays representing the cone positions of
        for all cone types
    """
    if isinstance(filename, str):
        filename = Path(filename)
    assert filename.is_file(), filename
    assert filename.suffix == ".lyt", filename
    data = filename.read_bytes()
    header_data, blocks_data = split_header_blocks(data)
    verify_lyt_header(header_data)

    all_cones_per_type = extract_cone_lists(blocks_data)

    all_cones_per_type_arrays = [
        np.array(cone_list) for cone_list in all_cones_per_type
    ]

    return all_cones_per_type_arrays

def generate_launch_description():

    # Include the robot_state_publisher launch file, provided by our own package. Force sim time to be enabled
    # !!! MAKE SURE YOU SET THE PACKAGE NAME CORRECTLY !!!

    package_name='dcaiti_control' #<--- CHANGE ME

    use_ros2_control = LaunchConfiguration('use_ros2_control')
    track = LaunchConfiguration('track')

    rsp = IncludeLaunchDescription(
                PythonLaunchDescriptionSource([os.path.join(
                    get_package_share_directory(package_name),'launch','rsp.launch.py'
                )]), launch_arguments={'use_sim_time': 'true', 'use_ros2_control': use_ros2_control}.items()
    )
    base_path = Path(get_package_share_directory(package_name))
    world_path = base_path / 'worlds'
    description_path = base_path  /'description'
    
    # Include the Gazebo launch file, provided by the gazebo_ros package

    gazebo = IncludeLaunchDescription(
                PythonLaunchDescriptionSource([os.path.join(
                    get_package_share_directory('ros_gz_sim'), 'launch', 'gz_sim.launch.py')]),
                    launch_arguments=[('gz_args', [f" -r -v 1 {world_path}/empty.sdf"])],
             )

    # Run the spawner node from the gazebo_ros package. The entity name doesn't really matter if you only have a single robot.
    spawn_entity = Node(package='ros_gz_sim', executable='create',
        arguments=[
            '-topic', 'robot_description',
            '-name', 'spawn_robot',
            '-z', '0.5',
            '-x', '-7',
            ],
        output='screen'
    )

    cone_positions = [x.reshape(-1,2) for x in load_lyt_file(base_path / 'tracks' / 'AU2_skidpad.lyt')]
    center = np.mean(cone_positions[ConeTypes.ORANGE_BIG], axis=0)

    cone_positions_centered = [x - center for x in cone_positions]
    

    cones_dict = {
        "unknown": cone_positions_centered[ConeTypes.UNKNOWN].tolist(),
        "yellow": cone_positions_centered[ConeTypes.YELLOW].tolist(),
        "blue": cone_positions_centered[ConeTypes.BLUE].tolist(),
        "orange_small": cone_positions_centered[ConeTypes.ORANGE_SMALL].tolist(),
        "orange_big": cone_positions_centered[ConeTypes.ORANGE_BIG].tolist(),
    }

    yellow_cone_xacro = f'{description_path}/yellow_cone.xacro'
    yellow_cone_config = Command(['xacro ', yellow_cone_xacro])
    blue_cone_xacro = f'{description_path}/blue_cone.xacro'
    blue_cone_config = Command(['xacro ', blue_cone_xacro])
    cone_spawner = []
    for i, (x,y) in enumerate(cones_dict['yellow']):
        spawn_cone = Node(package='ros_gz_sim', executable='create',
            arguments=[
                '-string', yellow_cone_config,
                '-name', f'yellow_cone_{i}',
                '-z', '0.3',
                '-y', str(y),
                '-x', str(x),
                ],
            output='screen'
        )
        cone_spawner.append(spawn_cone)
    for i, (x,y) in enumerate(cones_dict['blue']):
        spawn_cone = Node(package='ros_gz_sim', executable='create',
            arguments=[
                '-string', blue_cone_config,
                '-name', f'blue_cone_{i}',
                '-z', '0.3',
                '-y', str(y),
                '-x', str(x),
                ],
            output='screen'
        )
        cone_spawner.append(spawn_cone)

    # Bridge
    bridge = Node(
        package='ros_gz_bridge',
        executable='parameter_bridge',
        arguments=[
            '/clock@rosgraph_msgs/msg/Clock[ignition.msgs.Clock',
            '/boxes@vision_msgs/msg/Detection2DArray@ignition.msgs.AnnotatedAxisAligned2DBox_V',
            '/boxes_image@sensor_msgs/msg/Image@ignition.msgs.Image',
            '/camera_info@sensor_msgs/msg/CameraInfo@ignition.msgs.CameraInfo',
            '/lidar@sensor_msgs/msg/LaserScan@ignition.msgs.LaserScan',
            '/lidar/points@sensor_msgs/msg/PointCloud2@ignition.msgs.PointCloudPacked'
        ],
        output='screen'
    )
    
    diff_drive_spawner = Node(
        package="controller_manager",
        executable="spawner",
        arguments=["ackermann_steering_controller"],
    )

    joint_broad_spawner = Node(
        package="controller_manager",
        executable="spawner",
        arguments=["joint_broad"],
    )

    delayed_joint_broad_spawner = RegisterEventHandler(
      event_handler=OnProcessExit(
        target_action=spawn_entity,
        on_exit=[joint_broad_spawner],
      )
    )

    delayed_diff_drive_spawner = RegisterEventHandler(
      event_handler=OnProcessExit(
        target_action=spawn_entity,
        on_exit=[diff_drive_spawner],
      )
    )


    # Launch them all!
    return LaunchDescription([
        DeclareLaunchArgument(
            'use_ros2_control',
            default_value='true',
            description='Use ros2_control if true'),
        DeclareLaunchArgument(
            'track',
            default_value='AU2_skidpad.lyt',
            description='Which track to load'),
        bridge,
        rsp,
        gazebo,
        spawn_entity,
        delayed_diff_drive_spawner,
        delayed_joint_broad_spawner,
    ]+ cone_spawner ) 
