"""
I/O functions for the SD card
"""
from typing import Union, BinaryIO, Optional, Tuple, List, Literal, overload
from pathlib import Path
import cv2
import warnings

import numpy as np
from tqdm import tqdm

from miniscope_io.sdcard import SDLayout, SDConfig, DataHeader
from miniscope_io.exceptions import InvalidSDException, EndOfRecordingException
from miniscope_io.data import Frame


class SDCard:
    """
    I/O for data on an SDCard

    an instance of :class:`.sdcard.SDLayout` (typically in :mod:`.formats` ) configures how
    the data is laid out on the SD card. This class makes the i/o operations abstract over
    multiple layouts

    Args:
        drive (str, :class:`pathlib.Path`): Path to the SD card drive
        layout (:class:`.sdcard.SDLayout`): A layout configuration for an SD card

    """

    def __init__(
            self,
            drive: Union[str, Path],
            layout: SDLayout
        ):

        self.drive = drive
        self.layout = layout

        # Private attributes used when the file reading context is entered
        self._config = None  # type: Optional[SDConfig]
        self._f = None  # type: Optional[BinaryIO]
        self._frame = None  # type: Optional[int]
        self._frame_count = None # type: Optional[int]
        self._array = None  # type: Optional[np.ndarray]
        """
        n_pix x 1 array used to store pixels while reading buffers
        """
        self.positions = {}
        """
        A mapping between frame number and byte position in the video that makes for faster seeking :)
        
        As we read, we store the locations of each frame before reading it. Later, we can assign to
        `frame` to seek back to those positions. Assigning to `frame` works without caching position, but
        has to manually iterate through each frame.
        """



        # this doesn't seem to be true anymore? but the WRITE_KEYs are still in
        # both the firmware and reading code, just unused?
        #if not self.check_valid():
        #    raise InvalidSDException(f"The SD card at path {str(self.drive)} does not have the correct WRITEKEYs in its header!")

    # --------------------------------------------------
    # Properties
    # --------------------------------------------------

    @property
    def config(self) -> SDConfig:
        if self._config is None:
            with open(self.drive, 'rb') as sd:
                sd.seek(self.layout.sectors.config_pos, 0)
                configSectorData = np.frombuffer(sd.read(self.layout.sectors.size), dtype=np.uint32)

            self._config = SDConfig(
                **{
                    k: configSectorData[v]
                    for k, v in self.layout.config.dict().items()
                    if v is not None
                }
            )

        return self._config

    @property
    def position(self) -> Optional[int]:
        """
        When entered as context manager, the current position of the internal file descriptor
        """
        if self._f is None:
            return None

        return self._f.tell()


    @property
    def frame(self) -> Optional[int]:
        """
        When reading, the number of the frame that would be read if we were to call :meth:`.read`
        """
        if self._f is None:
            return None

        return self._frame

    @frame.setter
    def frame(self, frame:int):
        """
        Seek to a specific frame

        Arguments:
            frame (int): The frame to seek to!
        """
        if self._f is None:
            raise RuntimeError("Havent entered context manager yet! Cant change position without that!")

        if frame == self.frame:
            return

        if frame in self.positions.keys():
            self._f.seek(self.positions[frame], 0)
            self._frame = frame
            return
        else:
            # TODO find the nearest position we do have
            pass

        if frame < self.frame:
            # hard to go back, esp if we haven't already been here (we should have stashed the position)
            # just go to start of data and seek like normally (next case)
            self._f.seek(self.layout.sectors.data_pos, 0)
            self._frame = 0

        if frame > self.frame:
            for i in range(frame - self.frame):
                self.skip()

    @property
    def frame_count(self) -> int:
        """
        Total number of frames in recording.

        Inferred from :class:`~.sdcard.SDConfig.n_buffers_recorded` and
        reading a single frame to get the number of buffers per frame.
        """
        if self._frame_count is None:
            if self._f is None:
                with self as self_open:
                    frame = self_open.read(return_header=True)
                    headers = frame.headers

            else:
                # If we're already open, great, just return to the last frame
                last_frame = self.frame
                frame = self.read(return_header=True)
                headers = frame.headers
                self.frame = last_frame

            self._frame_count = int(np.ceil(
                (self.config.n_buffers_recorded + self.config.n_buffers_dropped) / len(headers)
            ))

        # if we have since read more frames than should be there, we update the frame count with a warning
        max_pos = np.max(list(self.positions.keys()))
        if max_pos > self._frame_count:
            warnings.warn(f'Got more frames than indicated in card header, expected {self._frame_count} but got {max_pos}')
            self._frame_count = int(max_pos)

        return self._frame_count



    # --------------------------------------------------
    # Context Manager methods
    # --------------------------------------------------

    def __enter__(self) -> 'SDCard':
        if self._f is not None:
            raise RuntimeError("Cant enter context, and open the file twice!")

        # init private attrs
        # create an empty frame to hold our data!
        self._array = np.zeros(
            (self.config.width * self.config.height, 1),
            dtype=np.uint8
        )
        self._pixel_count = 0
        self._last_buffer_n = 0
        self._frame = 0

        self._f = open(self.drive, 'rb')
        # seek to the start of the data
        self._f.seek(self.layout.sectors.data_pos, 0)
        # store the 0th frame position
        self.positions[0] = self.layout.sectors.data_pos

        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._f.close()
        self._f = None
        self._frame = 0

    # --------------------------------------------------
    # read methods
    # --------------------------------------------------
    def _read_data_header(self, sd:BinaryIO) -> DataHeader:
        """
        Given an already open file buffer opened in bytes mode,
         seeked to the start of a frame, read the data header
        """

        # read one word first, I think to get the size of the rest of the header,
        # that sort of breaks the abstraction (it assumes the buffer length is always at position 0)
        # but we'll roll with it for now
        dataHeader = np.frombuffer(sd.read(self.layout.word_size), dtype=np.uint32)
        dataHeader = np.append(
            dataHeader,
            np.frombuffer(
                sd.read((dataHeader[self.layout.buffer.length] - 1) * self.layout.word_size),
                dtype=np.uint32
            )
        )
        # use construct because we're already sure these are ints from the numpy casting
        #https://docs.pydantic.dev/latest/usage/models/#creating-models-without-validation
        header = DataHeader.model_construct(
            **{
                k: dataHeader[v]
                for k, v in self.layout.buffer.model_dump().items()
                if v is not None
            })
        return header

    def _n_frame_blocks(self, header: DataHeader) -> int:
        """
        Compute the number of blocks for a given frame buffer

        Not sure how this works!
        """
        n_blocks = int(
            (
                header.data_length +
                (header.length * self.layout.word_size) +
                (self.layout.sectors.size - 1)
            ) / self.layout.sectors.size
        )
        return n_blocks

    def _read_size(self, header: DataHeader) -> int:
        """
        Compute the number of bytes to read for a given buffer

        Not sure how this works with :meth:`._n_frame_blocks`, but keeping
        them separate in case they are separable actions for now
        """
        n_blocks = self._n_frame_blocks(header)
        read_size = (n_blocks * self.layout.sectors.size) - \
            (header.length * self.layout.word_size)
        return read_size

    def _read_buffer(self, sd: BinaryIO, header: DataHeader) -> np.ndarray:
        """
        Read a single buffer from a frame.

        Each frame has several buffers, so for a given frame we read them until we get another that's zero!
        """
        data = np.frombuffer(
            sd.read(self._read_size(header)),
            dtype=np.uint8
        )
        return data

    def _trim(self, data:np.ndarray, expected_size:int) -> np.ndarray:
        """
        Trim or pad an array to match an expected size
        """
        if data.shape[0] != expected_size:
            warnings.warn(
                f"Expected buffer data length: {expected_size}, got data with shape {data.shape}. Padding to expected length")

            # trim if too long
            if data.shape[0] > expected_size:
                data = data[0:expected_size]
            # pad if too short
            else:
                data = np.pad(data, (0, expected_size - data.shape[0]))

        return data

    @overload
    def read(self, return_header: Literal[True]= True) -> Frame: ...

    @overload
    def read(self, return_header: Literal[False] = False) -> np.ndarray: ...

    def read(
            self,
            return_header:bool = False
         ) -> Union[np.ndarray, Frame]:
        """
        Read a single frame

        Arguments:
            return_header (bool): If `True`, return headers from individual buffers (default `False`)

        Return:
            :class:`numpy.ndarray` , or a tuple(ndarray, List[:class:`~.DataHeader`]) if `return_header` is `True`
        """
        if self._f is None:
            raise RuntimeError('File is not open! Try entering the reader context by using it like `with sdcard:`')

        self._array[:] = 0
        pixel_count = 0
        last_buffer_n = 0
        headers = []
        while True:
            # stash position before reading header
            last_position = self._f.tell()
            try:
                header = self._read_data_header(self._f)
            except ValueError as e:
                if 'read length must be non-negative' in str(e):
                    # end of file! Value error thrown because the dataHeader will be blank,
                    # and thus have a value of 0 for the header size, and we can't read 0 from the card.
                    self._f.seek(last_position, 0)
                    raise EndOfRecordingException("Reached the end of the video!")
                else:
                    raise e
            except IndexError as e:
                if 'index 0 is out of bounds for axis 0 with size 0' in str(e):
                    # end of file if we are reading from a disk image without any additional space on disk
                    raise EndOfRecordingException("Reached the end of the video!")
                else:
                    raise e

            if header.frame_buffer_count == 0 and last_buffer_n > 0:
                # we are in the next frame!
                # rewind to the beginning of the header, and return
                # the last_position is the start of the header for this frame
                self._f.seek(last_position, 0)
                self._frame += 1
                self.positions[self._frame] = last_position
                frame = np.reshape(self._array, (self.config.width, self.config.height))
                if return_header:
                    return Frame.model_construct(frame=frame, headers=headers)
                else:
                    return frame

            # grab buffer data and stash
            headers.append(header)
            data = self._read_buffer(self._f, header)
            data = self._trim(data, header.data_length)
            self._array[pixel_count:pixel_count + header.data_length, 0] = data
            pixel_count += header.data_length
            last_buffer_n = header.frame_buffer_count

    # --------------------------------------------------
    # Write methods
    # --------------------------------------------------

    def to_video(
        self,
        path: Union[Path,str],
        fourcc: Literal['GREY', 'mp4v', 'XVID'] = 'GREY',
        isColor: bool = False,
        force: bool = False,
        progress: bool = True
    ):
        """
        Save contents of SD card to video with opencv

        Args:
            path (:class:`pathlib.Path`): Output video path, with video extension ``.avi`` or ``.mp4``
            fourcc (str): FourCC code used with opencv. Other codecs may be available depending on your opencv installation, but by default opencv supports one of:
                * ``GREY`` (default)
                * ``mp4v``
                * ``XVID``
            isColor (bool): Indicates whether output video is in color (default: `False`)
            force (bool): If `True`, overwrite output video if one already exists (default: `False`)
            progress (bool): If `True` (default) show progress bar.
        """
        path = Path(path)
        if path.exists() and not force:
            raise FileExistsError(f"{str(path)} already exists, not overwriting. Use force=True to overwrite.")

        if progress:
            pbar = tqdm(total=self.frame_count)

        writer = cv2.VideoWriter(
            str(path),
            cv2.VideoWriter_fourcc(*fourcc),
            self.config.fs,
            (self.config.width, self.config.height),
            isColor=isColor
        )

        # wrap in try block so we always close video writer class
        try:
            # Open video context manager
            with self:
                # catch the stop iteration signal
                try:
                    while True:
                        # this is sort of an awkward stack, should probably make a
                        # generator version of `read`
                        frame = self.read(return_header=False)
                        writer.write(frame)

                        if progress:
                            pbar.update()

                except StopIteration:
                    # end of the video!
                    pass

        finally:
            writer.release()
            if progress:
                pbar.close()

    # --------------------------------------------------
    # General Methods
    # --------------------------------------------------


    def skip(self):
        """
        Skip a frame

        Read the buffer headers to determine buffer sizes and just seek ahead
        """
        if self._f is None:
            raise RuntimeError('File is not open! Try entering the reader context by using it like `with sdcard:`')

        last_position = self._f.tell()
        header = self._read_data_header(self._f)

        if header.frame_buffer_count != 0:
            self._f.seek(last_position, 0)
            raise RuntimeError("Did not start at the first buffer of a frame! Something is wrong with the way seeking is working. Rewound to where we started")

        while True:
            # jump ahead according to the last header we read
            read_size = self._read_size(header)
            self._f.seek(read_size, 1)

            # stash position before reading next buffer header
            last_position = self._f.tell()
            header = self._read_data_header(self._f)

            # if the frame is over, return to the start of the buffer and break,
            # incrementing the current frame count
            if header.frame_buffer_count == 0:
                self._f.seek(last_position, 0)
                self._frame += 1
                self.positions[self.frame] = last_position
                break

    def check_valid(self) -> bool:
        """
        Checks that the header sector has the appropriate write keys in it

        Returns:
            bool - True if valid, False if not
        """
        with open(self.drive, 'rb') as sd:
            sd.seek(self.layout.sectors.header_pos, 0)
            headerSectorData = np.frombuffer(sd.read(self.layout.sectors.size), dtype=np.uint32)

            valid = False
            if headerSectorData[0] == self.layout.write_key0 and \
                    headerSectorData[1] == self.layout.write_key1 and \
                    headerSectorData[2] == self.layout.write_key2 and \
                    headerSectorData[3] == self.layout.write_key3:
                valid = True

        return valid
