from io import SEEK_CUR
import fractions, re
from struct import Struct
from collections import namedtuple
from itertools import accumulate

from ..utils import get_video_format, get_audio_format, stream_spec, get_samplesize
from .. import plugins

# https://docs.microsoft.com/en-us/previous-versions//dd183376(v=vs.85)?redirectedfrom=MSDN


class FlagProcessor:
    def __init__(self, name, flags, masks, defaults):
        self.template = namedtuple(
            name,
            flags,
            defaults=defaults,
        )
        self.masks = self.template._make(masks)

    def default(self):
        return self.template()

    def unpack(self, flags):
        return self.template._make((bool(flags & mask) for mask in self.masks))

    def pack(self, flags):
        return sum((mask if flag else 0 for flag, mask in zip(flags, self.masks)))


class StructProcessor:
    def __init__(self, name, format, fields, defaults=None, **flags):
        if "S" in format or "C" in format:
            # expand the format
            m = re.match(r"([<>!=])?(.+)", format)
            fmt_items = [
                (int(m[1]) if m[1] else 1, m[2])
                for m in re.finditer(r"(\d*)([xcCbB?hHiIlLqQnNefdsSpP])", m[2])
            ]
            fmt_counts = [1 if f in "sSp" else count for count, f in fmt_items]
            fmt_offsets = list((0, *accumulate(fmt_counts)))
            is_str = [False] * fmt_offsets[-1]
            for itm, offset in zip(fmt_items, fmt_offsets[:-1]):
                is_str[offset] = itm[1] in "SC"
            self.is_str = [fields[i] for i, tf in enumerate(is_str) if tf]
            format = format.replace("C", "c").replace("S", "s")
        else:
            self.is_str = ()

        self.struct = Struct(format)
        self.template = namedtuple(name, fields, defaults=defaults)
        self.flags = ((k, FlagProcessor(*v)) for k, v in flags.items())

    def default(self):
        data = self.template()
        return data._replace(**{k: proc.default() for k, proc in self.flags})

    def _unpack(self, data):
        data = self.template._make(data)
        return data._replace(
            **{field: getattr(data, field).decode("utf-8") for field in self.is_str},
            **{k: proc.unpack(getattr(data, k)) for k, proc in self.flags},
        )

    def unpack(self, buffer):
        return self._unpack(self.struct.unpack(buffer))

    def unpack_from(self, buffer, offset=0):
        return self._unpack(self.struct.unpack_from(buffer, offset))

    def _pack(self, ntuple):
        return ntuple._replace(
            **{k: proc.pack(getattr(ntuple, k)) for k, proc in self.flags},
            **{field: ntuple[field].encode("utf-8") for field in self.is_str},
        )

    def pack(self, ntuple):
        return self.struct.pack(*self._pack(ntuple))

    def pack_into(self, buffer, offset, ntuple):
        self.struct.pack_into(buffer, offset, *self._pack(ntuple))

    @property
    def size(self):
        return self.struct.size


AVIMainHeader = StructProcessor(
    "Avih",
    "<10I",
    (
        "micro_sec_per_frame",
        "max_bytes_per_sec",
        "padding_granularity",
        "flags",
        "total_frames",
        "initial_frames",
        "streams",
        "suggested_buffer_size",
        "width",
        "height",
    ),
    (0,) * 10,
    flags=(
        "AvihFlags",
        (
            "copyrighted",
            "has_index",
            "is_interleaved",
            "must_use_index",
            "was_capture_file",
        ),
        (
            int("0x00020000", 0),
            int("0x00000010", 0),
            int("0x00000100", 0),
            int("0x00000020", 0),
            int("0x00010000", 0),
        ),
        (False,) * 5,
    ),
)


AVIStreamHeader = StructProcessor(
    "AVISTREAMHEADER",
    "<4S4SI2H8I4h",
    (
        "fcc_type",  # 'auds','mids','txts','vids'
        "fcc_handler",
        "flags",
        "priority",
        "language",
        "initial_frame",
        "scale",
        "rate",
        "start",
        "length",
        "suggested_buffer_size",
        "quality",
        "sample_size",
        "frame_left",
        "frame_top",
        "frame_right",
        "frame_bottom",
    ),
    (b"\0" * 4, b"\0" * 4, *((0,) * 15)),
    flags=(
        "StrhFlags",
        (
            "video_pal_changes",
            "disabled",
        ),
        (
            int("0x00000001", 0),
            int("0x00010000", 0),
        ),
        (False,) * 2,
    ),
)

# PCM audio
WAVE_FORMAT_PCM = 1
# IEEE floating-point audio
WAVE_FORMAT_IEEE_FLOAT = 3
WAVE_FORMAT_EXTENSIBLE = int("FFFE", 16)  # /* Microsoft, 65534 */

BitmapInfoHeader = StructProcessor(
    "BITMAPINFOHEADER",
    "IiiHH4sIiiII",
    (
        "size",
        "width",
        "height",
        "planes",
        "bit_count",
        "compression",  # convert to str if 1st byte is >=4
        "size_image",
        "x_pels_per_meter",
        "y_pels_per_meter",
        "clr_used",
        "clr_important",
    ),
    (0,) * 11,
)

WaveFormatEx = StructProcessor(
    "WAVEFORMATEX",
    "HHIIHH",
    (
        "format_tag",
        "channels",
        "samples_per_sec",
        "avg_bytes_per_sec",
        "block_align",
        "bits_per_sample",
    ),
    (0,) * 6,
)

WaveFormatExtensible = StructProcessor(
    "WAVEFORMATEXTENSIBLE",
    "HHIH14s",
    (
        "size",
        "samples",
        "channel_mask",
        "sub_format_wave",
        "sub_format_rest",
    ),
    (*((0,) * 3), 0, "\0" * 14),
)


VideoPropHeader = StructProcessor(
    "VPRP",
    "5IHH3I",
    (
        "video_format_token",
        "video_standard",
        "vertical_refresh_rate",
        "h_total_in_t",
        "v_total_in_lines",
        "frame_aspect_ratio_y",
        "frame_aspect_ratio_x",
        "frame_width_in_pixels",
        "frame_height_in_lines",
        "field_per_frame",
    ),
    ((0,) * 10),
)

VPRP_VideoField = StructProcessor(
    "VPRP_VIDEO_FIELD_DESC",
    "8I",
    (
        "compressed_bm_height",
        "compressed_bm_width",
        "valid_bm_height",
        "valid_bm_width",
        "valid_bm_x_offset",
        "valid_bm_y_offset",
        "video_x_offset_in_t",
        "video_y_valid_start_line",
    ),
    ((0,) * 8),
)


ChunkHeader = StructProcessor("CHDR", "<4SI", ("id", "datasize"))


fcc_types = dict(vids="v", auds="a", txts="s")  # , mids="midi")


def read_chunk_header(f):
    b = f.read(ChunkHeader.size)
    id, datasize = ChunkHeader.unpack(b)
    list_type = None
    if id in ("RIFF", "LIST"):
        list_type = f.read(4).decode("utf-8")
        datasize -= 4
    chunksize = datasize + 1 if datasize % 2 else datasize
    return id, datasize, chunksize, list_type


def get_chunk_header(b, offset=0):
    id, datasize = ChunkHeader.unpack_from(b, offset)
    offset += ChunkHeader.size
    list_type = None
    if id in ("RIFF", "LIST"):
        list_type = b[offset : offset + 4].decode("utf-8")
        offset += 4
        datasize -= 4
    chunksize = datasize + 1 if datasize % 2 else datasize
    return offset, chunksize, id, list_type


def get_stream_header(b, offset, end):
    data = {}

    offset, chunksize, id, _ = get_chunk_header(b, offset)
    data[id] = strh = AVIStreamHeader.unpack_from(b, offset)
    offset += chunksize

    offset, chunksize, id, _ = get_chunk_header(b, offset)
    if strh.fcc_type == "vids":
        data[id] = BitmapInfoHeader.unpack_from(b, offset)

        # if 1st byte is a readable ascii char
        compression = data[id].compression
        comp_val = compression[0]
        data[id] = data[id]._replace(
            compression=comp_val if comp_val < 32 else compression.decode("utf-8")
        )

        # offset += chunksize
        # while offset < end:
        #     offset, chunksize, id, _ = get_chunk_header(b, offset)
        #     if id == "vprp":
        #         vprp = VideoPropHeader.unpack_from(b, offset)
        #         offset += VideoPropHeader.size
        #         ninfo = VPRP_VideoField.size
        #         field_info = [
        #             VPRP_VideoField.unpack_from(b, i)
        #             for i in range(offset, offset + ninfo * vprp.field_per_frame, ninfo)
        #         ]
        #         data[id] = namedtuple(
        #             type(vprp).__name__, (*vprp._fields, "field_info")
        #         )(*vprp, field_info)
        #         break
        #     else:
        #         offset += chunksize

    elif strh.fcc_type == "auds":
        strf = WaveFormatEx.unpack_from(b, offset)
        if strf.format_tag == WAVE_FORMAT_EXTENSIBLE:
            strfext = WaveFormatExtensible.unpack_from(b, offset + WaveFormatEx.size)
            strf = namedtuple(
                type(strfext).__name__, (*strf._fields, *strfext._fields)
            )(strfext.sub_format_wave, *strf[1:], *strfext)
        data[id] = strf
    else:
        raise RuntimeError(f"Unsupported stream type: {strh.fcc_type}")

    return data


def _seek(f, n):
    try:
        f.seek(n, SEEK_CUR)
    except:
        f.read(n)


def read_header(f, pix_fmt=None):

    # read the RIFF header
    id, datasize, chunksize, list_type = read_chunk_header(f)
    if id != "RIFF" or list_type != "AVI ":
        raise RuntimeError(f"File stream is not AVI")

    # read the hdrl chunk
    id, datasize, chunksize, list_type = read_chunk_header(f)
    if id != "LIST" and list_type != "hdrl":
        raise RuntimeError(f"AVI is missing header chunk")
    b = f.read(datasize)
    if chunksize > datasize:
        _seek(f, 1)

    # read until encountering the movi list
    while True:
        id, _, chunksize, list_type = read_chunk_header(f)
        if list_type == "movi":
            break
        _seek(f, chunksize)

    # parse hdrl LIST chunk
    offset, chunksize, id, list_type = get_chunk_header(b)
    if id != "avih":
        raise RuntimeError("missing avi chunk")
    avih = AVIMainHeader.unpack_from(b, offset)
    offset += chunksize
    streams = []
    while True:
        try:
            offset, chunksize, id, list_type = get_chunk_header(b, offset)
        except:
            break
        if list_type != "strl":
            break

        streams.append(get_stream_header(b, offset, offset + chunksize))
        offset += chunksize

    def get_stream_info(i, strl, use_ya):
        strh = strl["strh"]
        strf = strl["strf"]
        type = fcc_types[strh.fcc_type]  # raises if not valid type
        info = dict(index=i, type=type)
        if type == fcc_types["vids"]:
            info["frame_rate"] = fractions.Fraction(strh.rate, strh.scale)
            info["width"] = strf.width
            info["height"] = abs(strf.height)
            bpp = strf.bit_count
            compression = strf.compression
            # force unsupported pixel formats
            info["pix_fmt"] = (
                {"Y800": "gray", "RGBA": "rgba"}.get(compression, None)
                if isinstance(compression, str)
                else (compression, bpp)
                if compression
                else "rgba64le"
                if bpp == 64
                else "rgb48le"
                if bpp == 48
                else ("ya16le" if use_ya else "grayf32le")
                if bpp == 32
                else "rgb24"
                if bpp == 24
                else ("ya8" if use_ya else "gray16le")
                if bpp == 16
                else None
            )
            # vprp = strl.get("vprp", None)
            # info["dar"] = (
            #     fractions.Fraction(vprp.frame_aspect_ratio_x, vprp.frame_aspect_ratio_y)
            #     if vprp
            #     else None
            # )
            info["dtype"], info["shape"] = get_video_format(
                info["pix_fmt"], (info["width"], info["height"])
            )
        elif type == fcc_types["auds"]:  #'audio'
            info["sample_rate"] = strf.samples_per_sec
            info["channels"] = strf.channels

            strf_format = (
                strf.format_tag,
                strf.bits_per_sample,
            )

            info["sample_fmt"] = {
                (WAVE_FORMAT_PCM, 8): "u8",
                (WAVE_FORMAT_PCM, 16): "s16",
                (WAVE_FORMAT_PCM, 32): "s32",
                (WAVE_FORMAT_PCM, 64): "s64",
                (WAVE_FORMAT_IEEE_FLOAT, 32): "flt",
                (WAVE_FORMAT_IEEE_FLOAT, 64): "dbl",
            }.get(strf_format, strf_format)
            # TODO: if need arises, resolve more formats, need to include codec names though
            info["dtype"], info["shape"] = get_audio_format(
                info["sample_fmt"], info["channels"]
            )
        return info

    return [get_stream_info(i, strl, pix_fmt) for i, strl in enumerate(streams)], (
        avih,
        streams,
    )


re_movi = re.compile(r"\d{2}(?:wb|db|dc|tx)")


def read_frame(f):
    while True:
        id, datasize, chunksize, list_type = read_chunk_header(f)
        if not list_type:
            m = re_movi.match(id)
            if m:  # data chunk found
                b = f.read(datasize)
                if chunksize > datasize:
                    _seek(f, chunksize - datasize)
                return int(id[:2]), b
            else:
                _seek(f, chunksize)

        id, datasize, chunksize, list_type = read_chunk_header(f)


#######################################################################################################


class AviReader:
    def __init__(self):
        self._f = None
        self.ready = False  #:bool: True if AVI headers has been processed
        self.streams = None  #:dict: Stream headers keyed by stream id (int key)
        self.itemsizes = None  #:dict: sample size of each stream in bytes

        hook = plugins.get_hook()
        self.converters = {"v": hook.bytes_to_video, "a": hook.bytes_to_audio}
        #:dict : bytes to media data object conversion functions keyed by stream type

    def start(self, f, pix_fmt=None):
        self._f = f
        hdr = read_header(self._f, pix_fmt)[0]

        cnt = {"v": 0, "a": 0, "s": 0}

        def set_stream_info(hdr):
            st_type = hdr["type"]
            id = cnt[st_type]
            cnt[st_type] += 1
            return {
                "spec": stream_spec(id, st_type),
                **hdr,
            }

        self.streams = {v["index"]: set_stream_info(v) for v in hdr}
        self.itemsizes = {
            v["index"]: get_samplesize(v["shape"], v["dtype"]) for v in hdr
        }
        self.ready = True

    def __next__(self):
        i = d = None
        while i is None:  # None if unknown frame format, skip
            try:
                i, d = read_frame(self._f)
            except:
                raise StopIteration
        return i, d

    def __iter__(self):
        return self

    def from_bytes(self, id, b):
        info = self.streams[id]
        return self.converters[info["type"]](
            b=b, dtype=info["dtype"], shape=info["shape"], squeeze=False
        )


# (
#     "hdrl",
#     [
#         (
#             "avih",
#             {
#                 "micro_sec_per_frame": 66733,
#                 "max_bytes_per_sec": 3974198,
#                 "padding_granularity": 0,
#                 "flags": 0,
#                 "total_frames": 0,
#                 "initial_frames": 0,
#                 "streams": 2,
#                 "suggested_buffer_size": 1048576,
#                 "width": 352,
#                 "height": 240,
#             },
#         ),
#         (
#             "strl",
#             [
#                 (
#                     "strh",
#                     {
#                         "fcc_type": "vids",
#                         "fcc_handler": "\x00\x00\x00\x00",
#                         "flags": 0,
#                         "priority": 0,
#                         "language": 0,
#                         "initial_frames": 0,
#                         "scale": 200,
#                         "rate": 2997,
#                         "start": 0,
#                         "length": 1073741824,
#                         "suggested_buffer_size": 1048576,
#                         "quality": 4294967295,
#                         "sample_size": 0,
#                         "frame_left": 0,
#                         "frame_top": 0,
#                         "frame_right": 352,
#                         "frame_bottom": 240,
#                     },
#                 ),
#                 (
#                     "strf",
#                     {
#                         "size": 40,
#                         "width": 352,
#                         "height": -240,
#                         "planes": 1,
#                         "bit_count": 24,
#                         "compression": "rgb24",
#                         "size_image": 253440,
#                         "x_pels_per_meter": 0,
#                         "y_pels_per_meter": 0,
#                         "clr_used": 0,
#                         "clr_important": 0,
#                     },
#                 ),
#                 (
#                     "vprp",
#                     {
#                         "video_format_token": 0,
#                         "video_standard": 0,
#                         "vertical_refresh_rate": 15,
#                         "h_total_in_t": 352,
#                         "v_total_in_lines": 240,
#                         "frame_aspect_ratio": Fraction(15, 22),
#                         "frame_width_in_pixels": 352,
#                         "frame_height_in_lines": 240,
#                         "field_per_frame": 1,
#                         "field_info": (
#                             {
#                                 "compressed_bm_height": 240,
#                                 "compressed_bm_width": 352,
#                                 "valid_bm_height": 240,
#                                 "valid_bm_width": 352,
#                                 "valid_bmx_offset": 0,
#                                 "valid_bmy_offset": 0,
#                                 "video_x_offset_in_t": 0,
#                                 "video_y_valid_start_line": 0,
#                             },
#                         ),
#                     },
#                 ),
#             ],
#         ),
#         (
#             "strl",
#             [
#                 (
#                     "strh",
#                     {
#                         "fcc_type": "auds",
#                         "fcc_handler": "\x01\x00\x00\x00",
#                         "flags": 0,
#                         "priority": 0,
#                         "language": 0,
#                         "initial_frames": 0,
#                         "scale": 1,
#                         "rate": 44100,
#                         "start": 0,
#                         "length": 1073741824,
#                         "suggested_buffer_size": 12288,
#                         "quality": 4294967295,
#                         "sample_size": 4,
#                         "frame_left": 0,
#                         "frame_top": 0,
#                         "frame_right": 0,
#                         "frame_bottom": 0,
#                     },
#                 ),
#                 (
#                     "strf",
#                     {
#                         "format_tag": 1,
#                         "channels": 2,
#                         "samples_per_sec": 44100,
#                         "avg_bytes_per_sec": 176400,
#                         "block_align": 4,
#                         "bits_per_sample": 16,
#                     },
#                 ),
#             ],
#         ),
#     ],
#     368,
# )