from __future__ import annotations

import logging
from pathlib import Path
from typing import Optional, List, Dict, Iterator

import cv2
import numpy as np
import open3d as o3d
import pandas as pd
from rosbags.highlevel import AnyReader
from rosbags.typesys import get_typestore, Stores

from bagsy.std import keep_only_alpha_numerical

logger = logging.getLogger(__name__)


# ================================= GENERIC CONVERTER =================================


class Attribute:

    def __init__(self, name: str):
        self.name = name
        self.path = []

    def append_path(self, p) -> Attribute:
        self.path = [p] + self.path
        return self

    def get_fullname(self):
        return '.'.join(self.path + [self.name])

    def get_value(self, msg: object):
        for p in self.path:
            msg = getattr(msg, p)

        if type(msg) == np.ndarray:
            return msg[int(self.name)]
        return getattr(msg, self.name)

    @staticmethod
    def from_path(p):
        path = p.split(".")
        a = Attribute(path[-1])
        a.path = path[:-1]
        return a

    def __str__(self):
        return self.get_fullname()

    def __repr__(self):
        return str(self)


def find_msg_attributes(msg: object) -> List[Attribute]:
    if type(msg) in [str, int, float, bool]:
        return []
    if type(msg) == np.ndarray:
        if len(msg.shape) == 1 and msg.shape[0] < 100:
            return [Attribute(str(i)) for i in range(msg.shape[0])]
        else:
            raise NotImplementedError()

    result: List[Attribute] = []
    for attribute in dir(msg):
        if callable(getattr(msg, attribute)) or attribute.startswith('__'):
            continue
        subattributes = find_msg_attributes(getattr(msg, attribute))
        if len(subattributes) > 0:
            result += [subattribute.append_path(attribute) for subattribute in subattributes]
        else:
            result.append(Attribute(attribute))
    return result


# ================================= SPECIFIC CONVERTERS =================================

def get_attributes_tf(_: object) -> List[Attribute | str]:
    return [Attribute.from_path(p) for p in ['child_frame_id',
                                             'header.frame_id', 'header.seq', 'header.stamp.nanosec',
                                             'header.stamp.sec',
                                             'transform.rotation.w', 'transform.rotation.x', 'transform.rotation.y',
                                             'transform.rotation.z',
                                             'transform.translation.x', 'transform.translation.y',
                                             'transform.translation.z']]


def parse_tf(output_folder: Path, topic: str, bag_timestamp: int,
             msg: object, attributes: List[Attribute | str]) -> Iterator[List]:
    for transform in msg.transforms:
        yield [attr.get_value(transform) for attr in attributes]


def get_attributes_image(_: object) -> List[Attribute | str]:
    return [Attribute.from_path(p) for p in ['encoding',
                                             'header.frame_id', 'header.seq',
                                             'header.stamp.sec', 'header.stamp.nanosec',
                                             'height', 'width', 'is_bigendian', 'step']] + ['channel', 'path']


def parse_image(output_folder: Path, topic: str, bag_timestamp: int,
                msg: object, attributes: List[Attribute | str]) -> Iterator[List]:
    topic_output_folder = output_folder / keep_only_alpha_numerical(topic)
    topic_output_folder.mkdir(exist_ok=True)
    img = np.array(msg.data)
    img = img.reshape((msg.height, msg.width, img.shape[0] // msg.height // msg.width))
    img_path = topic_output_folder / (str(bag_timestamp) + ".png")
    cv2.imwrite(str(img_path), img)
    yield ([attr.get_value(msg) for attr in attributes if type(attr) == Attribute]
           + [img.shape[0] // msg.height // msg.width, img_path.relative_to(output_folder)])


def get_attributes_pc2(_: object) -> List[Attribute | str]:
    return [Attribute.from_path(p) for p in ['header.seq', 'header.frame_id',
                                             'header.stamp.sec', 'header.stamp.nanosec',
                                             'height', 'width', 'is_bigendian', 'is_dense', 'point_step',
                                             'row_step']] + ['path']


POINTFIELD_DTYPE_MAP = {1: np.int8, 2: np.uint8, 3: np.int16, 4: np.uint16,
                        5: np.int32, 6: np.uint32, 7: np.float32, 8: np.float64}


def parse_pc2(output_folder: Path, topic: str, bag_timestamp: int,
              msg: object, attributes: List[Attribute | str]) -> Iterator[List]:
    topic_output_folder = output_folder / keep_only_alpha_numerical(topic)
    topic_output_folder.mkdir(exist_ok=True)
    bag_path = topic_output_folder / (str(bag_timestamp) + ".pcd")

    # Decode raw buffer, respecting row_step padding
    num_points = msg.width * msg.height
    raw = (np.frombuffer(bytes(msg.data), dtype=np.uint8)
           .reshape(msg.height, msg.row_step)[:, :msg.width * msg.point_step]
           .reshape(num_points, msg.point_step))
    field_data = {field.name: np.frombuffer(
        raw[:, field.offset: field.offset + np.dtype(POINTFIELD_DTYPE_MAP[field.datatype]).itemsize].tobytes(),
        dtype=POINTFIELD_DTYPE_MAP[field.datatype])
        for field in msg.fields}

    pcd = o3d.t.geometry.PointCloud()
    for name, arr in field_data.items():
        pcd.point[name] = o3d.core.Tensor(arr.reshape(-1, 1))

    o3d.t.io.write_point_cloud(bag_path, pcd)
    yield ([attr.get_value(msg) for attr in attributes if type(attr) == Attribute]
           + [bag_path.relative_to(output_folder)])


CONVERTERS = {
    "tf2_msgs/msg/TFMessage": [get_attributes_tf, parse_tf],
    "sensor_msgs/msg/Image": [get_attributes_image, parse_image],
    "sensor_msgs/msg/PointCloud2": [get_attributes_pc2, parse_pc2],
}


# ================================= MAIN ROUTINE =================================


def parse_bag(bag_path: Path, topics: List[str], output_path: Path) -> Dict[str, pd.DataFrame]:
    """
    Parse a bag to pandas Dataframe
    :param bag_path: The path to the path
    :param topics: The topics to parse
    :param output_path: The place where to save certain special type of data
                        (that cannot or should not be in a Dataframe)
    :return: A dictionary of the data that was parsed for each topic
    """
    typestore = get_typestore(Stores.ROS1_NOETIC)
    with AnyReader([bag_path], default_typestore=typestore) as reader:
        types = {}
        for connection in reader.connections:
            if connection.topic in topics:
                types[connection.topic] = connection.msgtype
        logger.info(f"Found topics {', '.join([f'{topic} ({msgtype})' for topic, msgtype in types.items()])}.")

        # print unknown topics
        unknown = list(filter(lambda x: x not in types.keys(), topics))
        if len(unknown) > 0:
            logger.warning(f"Unknown topics: {', '.join(unknown)}. They will be ignored.")

        # several connections can have the same name, usually /tf is a good example
        connections = [c for c in reader.connections if c.topic in types.keys()]
        rows: Dict[str, List[object]] = {topic: [] for topic in types.keys()}
        attributes: Dict[str, Optional[List[Attribute | str]]] = {topic: None for topic in types.keys()}
        not_implemented: List[str] = []

        # Last benchmark (02/02/26) shows that :
        # - deserialization is between 25% and 50% of the duration
        # - attr.get_value is between 35% and 55% of the duration
        # - creating the DataFrame is about 20% of the duration
        #
        # concat/loc to creates dataframe is significantly (> 100 times) slower than using python array
        #
        # Also this can't be parallelized: rosbags lib doesn't seem to be thread safe (inherent to rosbag structure).
        for index, (connection, bag_timestamp, rawdata) in enumerate(reader.messages(connections=connections)):
            topics = connection.topic
            if topics in not_implemented:
                continue
            msgtype = connection.msgtype
            msg = reader.deserialize(rawdata, connection.msgtype)
            converter = CONVERTERS.get(msgtype, None)

            if attributes[topics] is None:
                try:
                    attributes[topics] = find_msg_attributes(msg) if converter is None else converter[0](msg)
                    logger.info(f"'{topics}': found {len(attributes[topics])} attributes.")
                except NotImplementedError:
                    logger.error(f"{msgtype} type is not supported. "
                                 f"Creating a specific converter may fix the issue.")
                    not_implemented.append(topics)
                    continue

            if converter is None:
                rows[topics].append([bag_timestamp * 1e-9] + [attr.get_value(msg) for attr in attributes[topics]])
            else:
                for row in converter[1](output_path, topics, bag_timestamp, msg, attributes[topics]):
                    rows[topics].append([bag_timestamp * 1e-9] + row)

        return {topic: pd.DataFrame(rows, columns=["msg_timestamp"] + [str(attr) for attr in attributes[topics]])
                for (topic, row) in rows.items()}
