import json
from pydantic import PositiveInt
from typing import Literal, Dict, List
from collections.abc import Mapping
from mcap.writer import Writer
from flatten_dict import flatten
from time import time_ns
from collections import defaultdict
from functools import partial
from pathlib import Path
from functools import cache
from shutil import rmtree
from mcap_data_loader.utils.av_coder import AvCoder, AvCoderConfig
from mcap_data_loader.utils.mcap_utils import McapTool, MediaType
from mcap_data_loader.serialization.flb import McapFlatBuffersWriter, FlatBuffersSchemas
from airbot_data_collection.common.samplers.basis import DataSampler, DataSamplerConfig


class McapDataSamplerConfig(DataSamplerConfig):
    """Configuration for MCAP data sampler."""

    initial_builder_size: PositiveInt = 1024 * 1024  # 1 MB
    """Initial size of the FlatBuffers builder."""
    video_save_to: Literal["file", "folder", "both"] = "file"
    """Where to save the video data: 'file' for MCAP attachment, 'folder' for separate folder, 'both' for both."""
    av_coder: AvCoderConfig = AvCoderConfig()
    """Configuration for the AV coder."""


class McapDataSampler(DataSampler):
    _info: Dict[str, Dict[str, str]]

    def __init__(self, config: McapDataSamplerConfig):
        self.config = config

    def on_configure(self):
        """Configure the mcap data sampler."""
        self._mf_writer = McapFlatBuffersWriter(self.config.initial_builder_size)
        self._coders: Dict[str, AvCoder] = defaultdict(
            partial(AvCoder, config=self.config.av_coder)
        )
        self._frame_stamp_factor = int(1e9 / self.config.av_coder.time_base)
        return True

    def _create_writer(self, path: Path) -> Writer:
        return Writer(str(path)), True

    def _get_video_dir(self, path: Path) -> Path:
        return path.parent / path.stem

    def compose_path(self, directory: Path, round: int) -> Path:
        path = directory / f"{round}.mcap"
        # unset here to ensure a fresh writer for each file
        # but the writer is finished in save()
        self._mf_writer.unset_writer()
        self._mf_writer.set_writer(*self._create_writer(path))
        for coder in self._coders.values():
            coder.reset()
        return path

    def update(self, data: dict):
        """Update the data with the latest frames."""
        # print(f"Updating data: {data.keys()}...")
        flag = None
        for key in tuple(data.keys()):
            if flag := self._is_save_h264(key):
                frame = data[key]
                self._coders[key].encode_frame(
                    frame["data"], frame["t"] // self._frame_stamp_factor
                )
            else:
                flag = self._add_messages(key, [data[key]], [data["log_stamps"]])
            if flag:
                data.pop(key)
        return data

    def save(self, path: Path, data: dict) -> str:
        """Save the data to a MCAP file."""
        writer = self._mf_writer.get_writer()
        mcap_tool = McapTool(writer)
        info = self._info.copy()
        # add metadata
        self.add_config_metadata(writer, self.config)
        # Handle system info safely
        # TODO: save system info to attachment?
        # TODO: should remap info keys?
        system_info = info.pop("system", {})
        for key, value in system_info.items():
            flattened_value = (
                flatten(value, "path")
                if isinstance(value, Mapping)
                else {"value": value}
            )
            # Convert all values to strings
            string_dict = {k: json.dumps(v) for k, v in flattened_value.items()}
            writer.add_metadata(key, string_dict)
        writer.add_attachment(
            time_ns(),
            time_ns(),
            "component_info",
            MediaType.APPLICATION_JSON,
            json.dumps(info).encode("utf-8"),
        )
        log_stamps = data.pop("log_stamps")
        mcap_tool.add_log_stamps_attachment(log_stamps)
        mcap_tool.add_topic_statistics_attachment(self._mf_writer.topic_statistics)
        for key, values in data.items():
            if not self._add_messages(key, values, log_stamps):
                self.get_logger().warning(f"Unknown data type for key: {key}")
        if self._coders:
            video_save_to = self.config.video_save_to
            video_dir = self._get_video_dir(path)
            video_path = None
            for key, coder in self._coders.items():
                key = self.config.key_remap(key)
                video_bytes = coder.end()
                if video_save_to in {"file", "both"}:
                    writer.add_attachment(
                        time_ns(), time_ns(), key, MediaType.VIDEO_MP4, video_bytes
                    )
                if video_save_to in {"folder", "both"}:
                    video_dir.mkdir(exist_ok=True)
                    video_path = (
                        video_dir / f"{key.removeprefix('/').replace('/', '.')}.mp4"
                    )
                    with open(video_path, "wb") as f:
                        f.write(video_bytes)
            if video_path is not None:
                self.get_logger().info(f"Saved videos to folder: {video_dir}")
        writer.finish()
        return path

    def remove(self, path):
        if self.config.video_save_to in {"folder", "both"}:
            video_dir = self._get_video_dir(path)
            self.get_logger().info(f"Removing video folder: {video_dir}")
            rmtree(video_dir, ignore_errors=True)
        return super().remove(path)

    def _add_messages(
        self, key: str, values: List[dict], log_stamps: List[float]
    ) -> FlatBuffersSchemas:
        # self.get_logger().info(f"Adding messages for key: {key}")
        schema_type = self._key_to_schema_type(key)
        if schema_type is not FlatBuffersSchemas.NONE:
            color_save_type = self.config.save_type.color
            if schema_type is FlatBuffersSchemas.COMPRESSED_IMAGE:
                kwargs = {"format": color_save_type, "frame_id": "airbot"}
            elif schema_type is FlatBuffersSchemas.RAW_IMAGE:
                kwargs = {"encoding": "", "frame_id": "airbot"}
            else:
                kwargs = {}
            if len(log_stamps) != len(values):
                raise ValueError(
                    f"Log stamps length ({len(log_stamps)}) must match data values length ({len(values)})."
                )
            key = self.config.key_remap(key)
            _ = [
                self._mf_writer.add_message(
                    schema_type, key, value["data"], value["t"], log_stamps[i], **kwargs
                )
                for i, value in enumerate(values)
            ]
        # else:
        #     self.get_logger().warning(f"Unknown data type for key: {key}")
        return schema_type

    @cache
    def _is_save_h264(self, key: str) -> bool:
        return "/color/" in key and self.config.save_type.color == "h264"

    @cache
    def _key_to_schema_type(self, key: str) -> FlatBuffersSchemas:
        is_color = "/color/" in key
        if is_color:
            save_type = self.config.save_type.color
            if save_type == "jpeg":
                return FlatBuffersSchemas.COMPRESSED_IMAGE
            elif save_type == "raw":
                return FlatBuffersSchemas.RAW_IMAGE
            else:
                return FlatBuffersSchemas.NONE
        is_depth = "depth" in key
        if is_depth:
            save_type = self.config.save_type.depth
            if save_type == "raw":
                return FlatBuffersSchemas.RAW_IMAGE
            raise NotImplementedError
        elif key == "log_stamps":
            return FlatBuffersSchemas.NONE
        return FlatBuffersSchemas.FLOAT_ARRAY

    @classmethod
    def add_config_metadata(cls, writer: Writer, config: McapDataSamplerConfig):
        config_dict = config.model_dump(mode="json")
        config_dict.pop("initial_builder_size")
        for key, value in config_dict.items():
            # Convert all values in dict to strings for MCAP metadata
            # MCAP add_metadata expects dict with string values
            if isinstance(value, dict):
                string_dict = {k: json.dumps(v) for k, v in value.items()}
            else:
                string_dict = {"value": json.dumps(value)}
            writer.add_metadata(name=key, data=string_dict)
