"""
This module provides classes to write data to shared memory using Cyclone DDS.

Classes:
    SMBufferWriteField: Manages a numpy array backed by shared memory for writing.
    SMWriter: Writes shared memory buffers using DDS for synchronization.
"""

import atexit
import time
from multiprocessing import shared_memory
from typing import Any, Dict, List, Tuple

import numpy as np
from cyclonedds.domain import DomainParticipant

from airo_ipc.cyclone_shm.idl.defaults.buffer_nr import BufferNrSample
from airo_ipc.cyclone_shm.idl_shared_memory.base_idl import BaseIdl
from airo_ipc.cyclone_shm.patterns.ddswriter import DDSWriter


class SMBufferWriteField:
    """
    Manages a numpy array backed by shared memory for writing.

    This class creates a shared memory segment and a numpy array
    that uses the shared memory as its buffer for writing data.
    """

    def __init__(self, name: str, shape: Tuple[int, ...], dtype: Any, nbytes: int):
        """
        Initialize the shared memory buffer field.

        Args:
            name (str): The name of the shared memory block.
            shape (tuple): The shape of the numpy array.
            dtype (numpy.dtype): The data type of the numpy array.
            nbytes (int): The number of bytes in the shared memory buffer.
        """

        try:
            self.shm = shared_memory.SharedMemory(create=True, size=nbytes, name=name)
        except FileExistsError:
            # this can occur if the process was killed and the shared memory was not cleaned up
            print(f"Shared memory file {name} exists. Will unlink and re-create it.")

            shm_old = shared_memory.SharedMemory(create=False, size=nbytes, name=name)
            shm_old.unlink()
            self.shm = shared_memory.SharedMemory(create=True, size=nbytes, name=name)
        # Create a new shared memory block with the given name and size
        # Create a numpy array that uses the shared memory buffer
        self.shared_array: np.ndarray = np.ndarray(
            shape, dtype=dtype, buffer=self.shm.buf
        )

    def stop(self) -> None:
        """Close and unlink the shared memory segment."""
        self.shm.close()
        self.shm.unlink()


class SMWriter:
    """
    Writes shared memory buffers using Cyclone DDS for synchronization.

    This class writes data to shared memory and publishes buffer numbers
    using DDS to synchronize readers.
    """

    def __init__(
        self,
        domain_participant: DomainParticipant,
        topic_name: str,
        idl_dataclass: BaseIdl,
        nr_of_buffers: int = 3,
    ):
        """
        Initialize the shared memory writer.

        Args:
            domain_participant (DomainParticipant): The DDS domain participant.
            topic_name (str): The base name of the topic.
            idl_dataclass (BaseIdl): The template defining the buffer structure.
            nr_of_buffers (int): The number of shared memory buffers per field. Raise this value to lower the chance of data races. Lower values reduce memory usage at the risk of data races. This value should be synchronized between writers and readers.
        """
        self.domain_participant = domain_participant
        self.topic_name = topic_name
        self.buffer_template = idl_dataclass
        self.nr_of_buffers = nr_of_buffers

        # Create a DDS writer for buffer numbers
        self.buffer_nr_writer = DDSWriter(
            domain_participant=domain_participant,
            topic_name=f"{topic_name}__buffer_nr",
            idl_dataclass=BufferNrSample,
        )

        # Create shared memory buffers
        self.buffers: List[Dict[str, SMBufferWriteField]] = self.__make_shared_memory()
        self.buffer_idx = 0

        atexit.register(self.stop)

    def stop(self) -> None:
        """Close and unlink all shared memory buffers owned by this writer."""
        for buffer in self.buffers:
            for bufferfield in buffer.values():
                try:
                    bufferfield.stop()
                except Exception:
                    pass
        atexit.unregister(self.stop)

    def __call__(self, msg: BaseIdl) -> None:
        """
        Write data to shared memory and publish buffer number via DDS.

        Args:
            msg (BaseIdl): The data to write to shared memory.
        """
        # Rotate to the next buffer index
        self.buffer_idx = (self.buffer_idx + 1) % self.nr_of_buffers

        buffer = self.buffers[self.buffer_idx]
        # Write each field to the shared memory buffer
        for key, bufferfield in buffer.items():
            # Copy data from the message to the shared memory array
            bufferfield.shared_array[:] = getattr(msg, key)[:]

        # Publish the buffer number and timestamp via DDS
        self.buffer_nr_writer(BufferNrSample(timestamp=time.time(), nr=self.buffer_idx))

    def __make_shared_memory(self) -> List[Dict[str, SMBufferWriteField]]:
        """
        Create shared memory buffers based on the buffer template.

        Returns:
            List[Dict[str, SMBufferWriteField]]: A list of dictionaries, each containing buffer fields.
        """
        # Initialize a list to hold buffers for each buffer index
        buffers: List[Dict[str, SMBufferWriteField]] = [
            {} for _ in range(self.nr_of_buffers)
        ]

        # Iterate over each field defined in the buffer template
        for name, shape, dtype, nbytes in self.buffer_template.get_fields():
            # For each buffer index, create a shared memory field
            for buffer_idx in range(self.nr_of_buffers):
                buffers[buffer_idx][name] = SMBufferWriteField(
                    f"{self.topic_name}.{name}.buffer_{buffer_idx}",
                    shape,
                    dtype,
                    nbytes,
                )

        return buffers
