"""
Reading Dolby Bitstream Metadata

Unless otherwise stated, all § references here are to
`EBU Tech 3285 Supplement 6`_.

.. _EBU Tech 3285 Supplement 6: https://tech.ebu.ch/docs/tech/tech3285s6.pdf
"""

from enum import IntEnum, Enum
from struct import unpack
from dataclasses import dataclass, asdict
from typing import List, Tuple, Any, Union

from io import BytesIO


class SegmentType(IntEnum):
    """
    Metadata segment type.
    """
    EndMarker = 0x0
    DolbyE = 0x1
    # Reserved2 = 0x2
    DolbyDigital = 0x3
    # Reserved4 = 0x4
    # Reserved5 = 0x5
    # Reserved6 = 0x6
    DolbyDigitalPlus = 0x7
    AudioInfo = 0x8
    DolbyAtmos = 0x9
    DolbyAtmosSupplemental = 0xa

    @classmethod
    def _missing_(cls, val):
        return val


@dataclass
class DolbyDigitalPlusMetadata:
    """
    *Dolby Digital Plus* is Dolby's brand for multichannel surround
    on discrete formats that aren't AC-3 (Dolby Digital) or Dolby E. This
    metadata segment is present in ADM wave files created with a Dolby Atmos
    Production Suite.

    Where an AC-3 bitstream can contain multiple programs, a Dolby Digital
    Plus bitstream will only contain one program.
    """

    class DownMixLevelToken(Enum):
        """
        A gain coefficient used in several metadata fields for downmix
        scenarios.
        """

        PLUS_3DB = 0b000
        "+3 dB"

        PLUS_1_5DB = 0b001
        "+1.5 dB"

        UNITY = 0b010
        "0dB"

        MINUS_1_5DB = 0b011
        "-1.5 dB"

        MINUS_3DB = 0b100
        "-3 dB"

        MINUS_4_5DB = 0b101
        "-4.5 dB"

        MINUS_6DB = 0b110
        "-6 dB"

        MUTE = 0b111
        "-∞ dB"

    class DolbySurroundEncodingMode(Enum):
        """
        Dolby surround endcoding mode.
        """
        RESERVED = 0b11
        IN_USE = 0b10
        NOT_IN_USE = 0b01
        NOT_INDICATED = 0b00

    class BitStreamMode(Enum):
        """
        Dolby Digital Plus `bsmod` field
        § 4.3.2.2
        """

        COMPLETE_MAIN = 0b000
        "main audio service: complete main"

        MUSIC_AND_EFFECTS = 0b001
        "main audio service: music and effects"

        VISUALLY_IMPAIRED = 0b010
        "associated service: visually impaired"

        HEARING_IMPAIRED = 0b011
        "associated service: hearing impaired"

        DIALOGUE_ONLY = 0b100
        "associated service: dialogue"

        COMMENTARY = 0b101
        "associated service: commentary"

        EMERGENCY = 0b110
        "associated service: emergency"

        VOICEOVER_KARAOKE = 0b111
        """
        associated service: voice over *OR* main audio service: karaoke.
        If `acmod` is `0b001` (mono 1/0), this is voice-over, otherwise it
        should be interpreted as karaoke.
        """

    class AudioCodingMode(Enum):
        """
        Dolby Digital Plus `acmod` field
        § 4.3.2.3
        """
        RESERVED = 0b000
        CH_ORD_1_0 = 0b001
        "Mono"
        CH_ORD_2_0 = 0b010
        "L/R stereo"
        CH_ORD_3_0 = 0b011
        "LCR stereo"
        CH_ORD_2_1 = 0b100
        "LR + mono surround"
        CH_ORD_3_1 = 0b101
        "LCR + mono surround"
        CH_ORD_2_2 = 0b110
        "LR + LR surround"
        CH_ORD_3_2 = 0b111
        "LCR + LR surround"

    class CenterDownMixLevel(Enum):
        """
        § 4.3.3.1
        """

        DOWN_3DB = 0b00
        "Attenuate 3 dB"

        DOWN_45DB = 0b01
        "Attenuate 4.5 dB"

        DOWN_6DB = 0b10
        "Attenuate 6 dB"

        RESERVED = 0b11

    class SurroundDownMixLevel(Enum):
        """
        Dolby Digital Plus `surmixlev` field
        § 4.3.3.2
        """
        DOWN_3DB = 0b00
        DOWN_6DB = 0b01
        MUTE = 0b10
        RESERVED = 0b11

    class LanguageCode(int):
        """
        § 4.3.4.1

        Per ATSC/A52 § 5.4.2.12, this is not in use and always 0xFF.
        """
        pass

    class MixLevel(int):
        """
        § 4.3.6.2
        """
        pass

    class DialnormLevel(int):
        """
        § 4.3.4.4
        """
        pass

    class RoomType(Enum):
        """
        `roomtyp` 4.3.6.3
        """
        NOT_INDICATED = 0b00
        LARGE_ROOM_X_CURVE = 0b01
        SMALL_ROOM_FLAT_CURVE = 0b10
        RESERVED = 0b11

    class PreferredDownMixMode(Enum):
        """
        Indicates the creating engineer's preference of what the receiver
        should downmix.
        § 4.3.8.1
        """
        NOT_INDICATED = 0b00
        PRO_LOGIC = 0b01
        STEREO = 0b10
        PRO_LOGIC_2 = 0b11

    class SurroundEXMode(IntEnum):
        """
        Dolby Surround-EX mode.
        `dsurexmod` § 4.3.9.1
        """
        NOT_INDICATED = 0b00
        NOT_SEX = 0b01
        SEX = 0b10
        PRO_LOGIC_2 = 0b11

    class HeadphoneMode(IntEnum):
        """
        `dheadphonmod` § 4.3.9.2
        """
        NOT_INDICATED = 0b00
        NOT_DOLBY_HEADPHONE = 0b01
        DOLBY_HEADPHONE = 0b10
        RESERVED = 0b11

    class ADConverterType(Enum):
        STANDARD = 0
        HDCD = 1

    class StreamDependency(Enum):
        """
        Encodes `ddplus_info1.stream_type` field § 4.3.12.1
        """

        INDEPENDENT = 0
        DEPENDENT = 1
        INDEPENDENT_FROM_DOLBY_DIGITAL = 2
        RESERVED = 3

    class RFCompressionProfile(Enum):
        """
        `compr1` RF compression profile
        § 4.3.10 (fig 42)
        """
        NONE = 0
        FILM_STANDARD = 1
        FILM_LIGHT = 2
        MUSIC_STANDARD = 3
        MUSIC_LIGHT = 4
        SPEECH = 5

    #: Program ID number, this identifies the program in a multi-program
    #: element. § 4.3.1
    program_id: int

    #: `True` if LFE is enabled. § 4.3.2.1
    lfe_on: bool

    #: The kind of service of this stream. `bsmod` § 4.3.2.2
    bitstream_mode: BitStreamMode

    #: Indicates which channels are in use. `acmod` § 4.3.2.3
    audio_coding_mode: AudioCodingMode

    #: When the front three channels are in use, gives the center
    #: downmix level. ``
    center_downmix_level: CenterDownMixLevel

    #: When the surround channels are in use, gives the surround
    #: downmix level.
    surround_downmix_level: SurroundDownMixLevel

    #: If the `acmod` is LR, this indicates if the channels
    #: are encoded in Dolby Surround.
    dolby_surround_encoded: DolbySurroundEncodingMode

    #: `True` if there is a langcode present in the metadata.
    langcode_present: bool

    #: `True` if this bitstream is copyrighted.
    copyright_bitstream: bool

    #: `True` if this bitstream is original.
    original_bitstream: bool

    dialnorm: DialnormLevel

    #: Language code
    langcode: int

    #: `True` if `mixlevel` and `roomtype` are valid
    prod_info_exists: bool

    #: Mix level
    mixlevel: MixLevel

    #: Room Type
    roomtype: RoomType

    #: LoRo preferred center downmix level
    loro_center_downmix_level: DownMixLevelToken

    #: LoRo preferred surround downmix level
    loro_surround_downmix_level: DownMixLevelToken

    #: Preferred downmix mode
    downmix_mode: PreferredDownMixMode

    #: LtRt preferred center downmix level
    ltrt_center_downmix_level: DownMixLevelToken

    #: LtRt preferred surround downmix level
    ltrt_surround_downmix_level: DownMixLevelToken

    #: Surround-EX mode
    surround_ex_mode: SurroundEXMode

    #: Dolby Headphone mode
    dolby_headphone_encoded: HeadphoneMode

    ad_converter_type: ADConverterType
    compression_profile: RFCompressionProfile
    dynamic_range: RFCompressionProfile

    #: Indicates if this stream can be decoded independently or not
    stream_dependency: StreamDependency

    #: Data rate of this bitstream in kilobits per second
    datarate_kbps: int

    @staticmethod
    def load(buffer: bytes):
        assert len(buffer) == 96, "Dolby Digital Plus segment incorrect size, "
        "expected 96 got %i" % len(buffer)

        def program_id(b) -> int:
            return b

        def program_info(b):
            return (b & 0x40) > 0, \
                DolbyDigitalPlusMetadata.BitStreamMode(b & 0x38 >> 3), \
                DolbyDigitalPlusMetadata.AudioCodingMode(b & 0x7)

        def ddplus_reserved1(_):
            pass

        def surround_config(b):
            return (
                DolbyDigitalPlusMetadata.CenterDownMixLevel(b & 0x30 >> 4),
                DolbyDigitalPlusMetadata.SurroundDownMixLevel(b & 0xc >> 2),
                DolbyDigitalPlusMetadata.DolbySurroundEncodingMode(b & 0x3)
                )

        def dialnorm_info(b):
            return (b & 0x80) > 0, b & 0x40 > 0, b & 0x20 > 0, \
                DolbyDigitalPlusMetadata.DialnormLevel(b & 0x1f)

        def langcod(b) -> int:
            return b

        def audio_prod_info(b):
            return (b & 0x80) > 0, \
                DolbyDigitalPlusMetadata.MixLevel(b & 0x7c >> 2), \
                DolbyDigitalPlusMetadata.RoomType(b & 0x3)

        # loro_center_downmix_level, loro_surround_downmix_level
        def ext_bsi1_word1(b):
            return DolbyDigitalPlusMetadata.DownMixLevelToken(b & 0x38 >> 3), \
                DolbyDigitalPlusMetadata.DownMixLevelToken(b & 0x7)

        # downmix_mode, ltrt_center_downmix_level, ltrt_surround_downmix_level
        def ext_bsi1_word2(b):
            return DolbyDigitalPlusMetadata\
                .PreferredDownMixMode(b & 0xC0 >> 6), \
                DolbyDigitalPlusMetadata.DownMixLevelToken(b & 0x38 >> 3), \
                DolbyDigitalPlusMetadata.DownMixLevelToken(b & 0x7)

        # surround_ex_mode, dolby_headphone_encoded, ad_converter_type
        def ext_bsi2_word1(b):
            return DolbyDigitalPlusMetadata.SurroundEXMode(b & 0x60 >> 5), \
                DolbyDigitalPlusMetadata.HeadphoneMode(b & 0x18 >> 3), \
                DolbyDigitalPlusMetadata.ADConverterType(b & 0x4 >> 2)

        def ddplus_reserved2(_):
            pass

        def compr1(b):
            return DolbyDigitalPlusMetadata.RFCompressionProfile(b)

        def dynrng1(b):
            DolbyDigitalPlusMetadata.RFCompressionProfile(b)

        def ddplus_reserved3(_):
            pass

        def ddplus_info1(b):
            return DolbyDigitalPlusMetadata.StreamDependency(b & 0xc >> 2)

        def ddplus_reserved4(_):
            pass

        def datarate(b) -> int:
            return unpack("<H", b)[0]

        def reserved(_):
            pass

        pid = program_id(buffer[0])
        lfe_on, bitstream_mode, audio_coding_mode = program_info(buffer[1])
        ddplus_reserved1(buffer[2:2])
        center_downmix_level, surround_downmix_level, \
            dolby_surround_encoded = surround_config(buffer[4])
        langcode_present, copyright_bitstream, original_bitstream, \
            dialnorm = dialnorm_info(buffer[5])
        langcode = langcod(buffer[6])
        prod_info_exists, mixlevel, roomtype = audio_prod_info(buffer[7])

        loro_center_downmix_level, \
            loro_surround_downmix_level = ext_bsi1_word1(buffer[8])
        downmix_mode, ltrt_center_downmix_level, \
            ltrt_surround_downmix_level = ext_bsi1_word2(buffer[9])
        surround_ex_mode, dolby_headphone_encoded, \
            ad_converter_type = ext_bsi2_word1(buffer[10])

        ddplus_reserved2(buffer[11:14])
        compression = compr1(buffer[14])
        dynamic_range = dynrng1(buffer[15])
        ddplus_reserved3(buffer[16:19])
        stream_info = ddplus_info1(buffer[19])
        ddplus_reserved4(buffer[20:25])
        data_rate = datarate(buffer[25:27])
        reserved(buffer[27:69])

        return DolbyDigitalPlusMetadata(
                 program_id=pid, lfe_on=lfe_on,
                 bitstream_mode=bitstream_mode,
                 audio_coding_mode=audio_coding_mode,
                 center_downmix_level=center_downmix_level,
                 surround_downmix_level=surround_downmix_level,
                 dolby_surround_encoded=dolby_surround_encoded,
                 langcode_present=langcode_present,
                 copyright_bitstream=copyright_bitstream,
                 original_bitstream=original_bitstream,
                 dialnorm=dialnorm,
                 langcode=langcode,
                 prod_info_exists=prod_info_exists,
                 mixlevel=mixlevel,
                 roomtype=roomtype,
                 loro_center_downmix_level=loro_center_downmix_level,
                 loro_surround_downmix_level=loro_surround_downmix_level,
                 downmix_mode=downmix_mode,
                 ltrt_center_downmix_level=ltrt_center_downmix_level,
                 ltrt_surround_downmix_level=ltrt_surround_downmix_level,
                 surround_ex_mode=surround_ex_mode,
                 dolby_headphone_encoded=dolby_headphone_encoded,
                 ad_converter_type=ad_converter_type,
                 compression_profile=compression,
                 dynamic_range=dynamic_range,
                 stream_dependency=stream_info,
                 datarate_kbps=data_rate)


@dataclass
class DolbyAtmosMetadata:
    """
    Dolby Atmos Metadata Segment

    https://github.com/DolbyLaboratories/dbmd-atmos-parser/
    """

    class WarpMode(Enum):
        NORMAL = 0x00
        WARPING = 0x01
        DOWNMIX_PLIIX = 0x02
        DOWNMIX_LORO = 0x03
        NOT_INDICATED = 0x04

    tool_name: str
    tool_version: Tuple[int, int, int]
    warp_mode: WarpMode

    SEGMENT_LENGTH = 248
    TOOL_NAME_LENGTH = 64

    @classmethod
    def load(cls, data: bytes):

        assert len(data) == cls.SEGMENT_LENGTH
        # (f"DolbyAtmosMetadata segment is incorrect length, "
        #  f"expected {cls.SEGMENT_LENGTH} actual was {len(data)}")

        h = BytesIO(data)

        h.seek(32, 1)
        toolname = h.read(cls.TOOL_NAME_LENGTH)
        toolname = unpack("%is" % cls.TOOL_NAME_LENGTH, toolname)[0]
        toolname = toolname.decode('utf-8').strip('\0')

        vers = h.read(3)
        major, minor, fix = unpack("BBB", vers)

        h.seek(53, 1)

        a_val = unpack("B", h.read(1))[0]
        warp_mode = a_val & 0x7

        return DolbyAtmosMetadata(tool_name=toolname,
                                  tool_version=(major, minor, fix),
                                  warp_mode=DolbyAtmosMetadata
                                  .WarpMode(warp_mode))


@dataclass
class DolbyAtmosSupplementalMetadata:
    """
    Dolby Atmos supplemental metadata segment.

    https://github.com/DolbyLaboratories/dbmd-atmos-parser/blob/
    master/dbmd_atmos_parse/src/dbmd_atmos_parse.c
    """

    class BinauralRenderMode(Enum):
        BYPASS = 0x00
        NEAR = 0x01
        FAR = 0x02
        MID = 0x03
        NOT_INDICATED = 0x04

    object_count: int
    render_modes: List['DolbyAtmosSupplementalMetadata.BinauralRenderMode']
    trim_modes: List[int]

    MAGIC = 0xf8726fbd
    TRIM_CONFIG_COUNT = 9

    @classmethod
    def load(cls, data: bytes):

        trim_modes = []
        render_modes = []

        h = BytesIO(data)
        magic = unpack("<I", h.read(4))[0]
        assert magic == cls.MAGIC, "Magic value was not found"

        object_count = unpack("<H", h.read(2))[0]

        h.read(1)  # skip 1

        for _ in range(cls.TRIM_CONFIG_COUNT):
            auto_trim = unpack("B", h.read(1))
            trim_modes.append(auto_trim)

            h.read(14)  # skip 14

        h.read(object_count)  # skip object_count bytes

        for _ in range(object_count):
            binaural_mode = unpack("B", h.read(1))[0]
            binaural_mode &= 0x7
            render_modes.append(binaural_mode)

        return DolbyAtmosSupplementalMetadata(object_count=object_count,
                                              render_modes=render_modes,
                                              trim_modes=trim_modes)


class WavDolbyMetadataReader:
    """
    Reads Dolby bitstream metadata.
    """

    #: List of the Dolby Metadata Segments.
    #:
    #: Each list entry is a tuple of `SegmentType`, a `bool`
    #: indicating if the segment's checksum was valid, and the
    #: segment's parsed dataclass (or a `bytes` array if it was
    #: not recognized).
    segment_list: List[Tuple[Union[SegmentType, int], bool, Any]]

    version: Tuple[int, int, int, int]

    @staticmethod
    def segment_checksum(bs: bytes, size: int):
        retval = size
        for b in bs:
            retval += int(b)
            retval &= 0xff

        retval = ((~retval) + 1) & 0xff

        return retval

    def __init__(self, dbmd_data):
        self.segment_list = []

        h = BytesIO(dbmd_data)

        v_vec = []
        for _ in range(4):
            b = h.read(1)
            v_vec.insert(0, unpack("B", b)[0])

        self.version = tuple(v_vec)

        while True:
            stype = SegmentType(unpack("B", h.read(1))[0])
            if stype == SegmentType.EndMarker:
                break
            else:
                seg_size = unpack("<H", h.read(2))[0]
                seg_payload = h.read(seg_size)
                expected_checksum = WavDolbyMetadataReader\
                    .segment_checksum(seg_payload, seg_size)
                checksum = unpack("B", h.read(1))[0]

                segment = seg_payload
                if stype == SegmentType.DolbyDigitalPlus:
                    segment = DolbyDigitalPlusMetadata.load(segment)
                elif stype == SegmentType.DolbyAtmos:
                    segment = DolbyAtmosMetadata.load(segment)
                elif stype == SegmentType.DolbyAtmosSupplemental:
                    segment = DolbyAtmosSupplementalMetadata.load(segment)

                self.segment_list\
                    .append((stype, checksum == expected_checksum, segment))

    def dolby_digital_plus(self) -> List[DolbyDigitalPlusMetadata]:
        """
        Every valid Dolby Digital Plus metadata segment in the file.
        """
        return [x[2] for x in self.segment_list
                if x[0] == SegmentType.DolbyDigitalPlus and x[1]]

    def dolby_atmos(self) -> List[DolbyAtmosMetadata]:
        """
        Every valid Dolby Atmos metadata segment in the file.
        """
        return [x[2] for x in self.segment_list
                if x[0] == SegmentType.DolbyAtmos and x[1]]

    def dolby_atmos_supplemental(self) -> List[DolbyAtmosSupplementalMetadata]:
        """
        Every valid Dolby Atmos Supplemental metadata segment in the file.
        """
        return [x[2] for x in self.segment_list
                if x[0] == SegmentType.DolbyAtmosSupplemental and x[1]]

    def to_dict(self) -> dict:

        ddp = map(lambda x: asdict(x), self.dolby_digital_plus())
        atmos = map(lambda x: asdict(x), self.dolby_atmos())
        # atmos_sup = map(lambda x: asdict(x), self.dolby_atmos_supplemental())

        return dict(dolby_digital_plus=list(ddp),
                    dolby_atmos=list(atmos))
