import time
import numpy as np
import pymap3d as pm
from pathlib import Path
from rosbags.highlevel import AnyReader
import autonomous_proto
from rosbags.rosbag2 import Writer
from rosbags.typesys import Stores, get_typestore
from rosbags.typesys.stores.ros2_humble import (
    std_msgs__msg__UInt8MultiArray as UInt8MultiArray,
    std_msgs__msg__MultiArrayLayout as MultiArrayLayout,
    std_msgs__msg__MultiArrayDimension as MultiArrayDimension,
)

class BagNav:
    def __init__(self,
                 origin_lat = 40,
                 origin_lon = 116,
                 origin_alt = 0,
                 bag_path: Path = None
                 ):
        self.origin_lat = origin_lat
        self.origin_lon = origin_lon
        self.origin_alt = origin_alt
        self.enus = []
        self.msgs: list[autonomous_proto.Navigation] = []
        if bag_path is None:
            return
        with AnyReader([bag_path]) as reader:
            connections = [c for c in reader.connections if 'navigation' in c.topic]
            for connection, timestamp, rawdata in reader.messages(connections=connections):
                msg = reader.deserialize(rawdata, connection.msgtype)
                if hasattr(msg, 'data'):
                    self.msgs.append(autonomous_proto.Navigation.FromString(bytes(msg.data)))
        if len(self.msgs) == 0:
            return
        self.origin_nav = self.msgs[0]
        self.origin_lat = self.origin_nav.position.lat
        self.origin_lon = self.origin_nav.position.lon
        self.origin_alt = self.origin_nav.position.alt
        for m in self.msgs:
            self.enus.append(pm.geodetic2enu(m.position.lat,
                                             m.position.lon,
                                             m.position.alt,
                                             self.origin_nav.position.lat,
                                             self.origin_nav.position.lon,
                                             self.origin_nav.position.alt))

    def save_to_bag(self, points: list[tuple[float, float]]):
        if len(points) == 0:
            return
        nav = autonomous_proto.Navigation()
        nav.header.info.count = 0
        nav.header.info.module_name = autonomous_proto.MessageInfoModuleNameValue.undefined
        nav.header.info.topic_name = autonomous_proto.MessageInfoTopicNameValue.undefined

        with Writer(Path('spline_' + time.strftime("%Y_%m_%d_%H_%M_%S")), version=8) as writer:
            type_store = get_typestore(Stores.ROS2_HUMBLE)
            topic = 'navigation'
            msg_type = UInt8MultiArray.__msgtype__
            connection = writer.add_connection(topic, msg_type, typestore=type_store)
            timestamp = int(time.time() * 1e9)
            for p in points:
                lat, lon, alt = pm.enu2geodetic(p[0], p[1], 0, self.origin_lat, self.origin_lon, self.origin_alt)
                nav.position.lat = lat
                nav.position.lon = lon
                nav.position.alt = alt
                timestamp += 20000000
                data = np.frombuffer(nav.SerializeToString(), dtype=np.uint8)
                dims = [MultiArrayDimension(label='', size=data.size, stride=1)]
                layout = MultiArrayLayout(dim=dims, data_offset=0)
                message = UInt8MultiArray(layout=layout, data=data)
                writer.write(connection, timestamp, type_store.serialize_cdr(message, msg_type))
