import io
import time
from binascii import crc32
from collections.abc import Iterable
from typing import Literal, Optional, Union, cast, overload

from typing_extensions import Self

from aiokafka.codec import (
    gzip_decode,
    has_gzip,
    has_lz4,
    has_snappy,
    has_zstd,
    lz4_decode,
    snappy_decode,
    zstd_decode,
)
from aiokafka.errors import UnsupportedCodecError

from .struct import Struct
from .types import Bytes, Int8, Int32, Int64, Schema, UInt32


class Message(Struct):
    # FIXME: override __eq__/__repr__ methods from Struct

    BASE_FIELDS = (
        ("crc", UInt32),
        ("magic", Int8),
        ("attributes", Int8),
    )
    MAGIC0_FIELDS = (
        ("key", Bytes),
        ("value", Bytes),
    )
    MAGIC1_FIELDS = (
        ("timestamp", Int64),
        ("key", Bytes),
        ("value", Bytes),
    )
    SCHEMAS = [
        Schema(
            *BASE_FIELDS,
            *MAGIC0_FIELDS,
        ),
        Schema(
            *BASE_FIELDS,
            *MAGIC1_FIELDS,
        ),
    ]
    SCHEMA = SCHEMAS[1]
    CODEC_MASK = 0x07
    CODEC_GZIP = 0x01
    CODEC_SNAPPY = 0x02
    CODEC_LZ4 = 0x03
    CODEC_ZSTD = 0x04
    TIMESTAMP_TYPE_MASK = 0x08
    HEADER_SIZE = (
        22  # crc(4), magic(1), attributes(1), timestamp(8), key+value size(4*2)
    )

    @overload
    def __init__(
        self,
        *,
        value: Optional[bytes],
        key: Optional[bytes],
        magic: Literal[0],
        attributes: int,
        crc: int,
    ) -> None: ...

    @overload
    def __init__(
        self,
        *,
        value: Optional[bytes],
        key: Optional[bytes],
        magic: Literal[1],
        attributes: int,
        crc: int,
        timestamp: int,
    ) -> None: ...

    def __init__(
        self,
        *,
        value: Optional[bytes],
        key: Optional[bytes],
        magic: Literal[0, 1],
        attributes: int,
        crc: int,
        timestamp: Optional[int] = None,
    ) -> None:
        assert value is None or isinstance(value, bytes), "value must be bytes"
        assert key is None or isinstance(key, bytes), "key must be bytes"
        assert magic > 0 or timestamp is None, "timestamp not supported in v0"

        # Default timestamp to now for v1 messages
        if magic > 0 and timestamp is None:
            timestamp = int(time.time() * 1000)
        self.timestamp = timestamp
        self.crc = crc
        self._validated_crc: Optional[int] = None
        self.magic = magic
        self.attributes = attributes
        self.key = key
        self.value = value

    @property
    def timestamp_type(self) -> Optional[Literal[0, 1]]:
        """0 for CreateTime; 1 for LogAppendTime; None if unsupported.

        Value is determined by broker; produced messages should always set to 0
        Requires Kafka >= 0.10 / message version >= 1
        """
        if self.magic == 0:
            return None
        elif self.attributes & self.TIMESTAMP_TYPE_MASK:
            return 1
        else:
            return 0

    def encode(self, recalc_crc: bool = True) -> bytes:
        version = self.magic
        if version == 1:
            message = Message.SCHEMAS[version].encode(
                (
                    self.crc,
                    self.magic,
                    self.attributes,
                    self.timestamp,
                    self.key,
                    self.value,
                )
            )
        elif version == 0:
            message = Message.SCHEMAS[version].encode(
                (self.crc, self.magic, self.attributes, self.key, self.value)
            )
        else:
            raise ValueError(f"Unrecognized message version: {version}")
        if not recalc_crc:
            return message
        self.crc = crc32(message[4:])
        crc_field = self.BASE_FIELDS[0][1]
        return crc_field.encode(self.crc) + message[4:]

    @classmethod
    def decode(cls, data: Union[io.BytesIO, bytes]) -> Self:
        _validated_crc: Optional[int] = None
        if isinstance(data, bytes):
            _validated_crc = crc32(data[4:])
            data = io.BytesIO(data)
        # Partial decode required to determine message version
        crc, magic, attributes = (
            cls.BASE_FIELDS[0][1].decode(data),
            cls.BASE_FIELDS[1][1].decode(data),
            cls.BASE_FIELDS[2][1].decode(data),
        )
        if magic == 1:
            magic = cast(Literal[1], magic)
            timestamp, key, value = (
                cls.MAGIC1_FIELDS[0][1].decode(data),
                cls.MAGIC1_FIELDS[1][1].decode(data),
                cls.MAGIC1_FIELDS[2][1].decode(data),
            )
            msg = cls(
                value=value,
                key=key,
                magic=magic,
                attributes=attributes,
                crc=crc,
                timestamp=timestamp,
            )
        elif magic == 0:
            magic = cast(Literal[0], magic)
            key, value = (
                cls.MAGIC0_FIELDS[0][1].decode(data),
                cls.MAGIC0_FIELDS[1][1].decode(data),
            )
            msg = cls(
                value=value,
                key=key,
                magic=magic,
                attributes=attributes,
                crc=crc,
            )
        else:
            raise ValueError(f"Unrecognized message version: {magic}")

        msg._validated_crc = _validated_crc
        return msg

    def validate_crc(self) -> bool:
        if self._validated_crc is None:
            raw_msg = self.encode(recalc_crc=False)
            self._validated_crc = crc32(raw_msg[4:])
        return self.crc == self._validated_crc

    def is_compressed(self) -> bool:
        return self.attributes & self.CODEC_MASK != 0

    def decompress(
        self,
    ) -> list[Union[tuple[int, int, "Message"], tuple[None, None, "PartialMessage"]]]:
        assert self.value is not None
        codec = self.attributes & self.CODEC_MASK
        assert codec in (
            self.CODEC_GZIP,
            self.CODEC_SNAPPY,
            self.CODEC_LZ4,
            self.CODEC_ZSTD,
        )
        if codec == self.CODEC_GZIP:
            assert has_gzip(), "Gzip decompression unsupported"
            raw_bytes = gzip_decode(self.value)
        elif codec == self.CODEC_SNAPPY:
            assert has_snappy(), "Snappy decompression unsupported"
            raw_bytes = snappy_decode(self.value)
        elif codec == self.CODEC_LZ4:
            assert has_lz4(), "LZ4 decompression unsupported"
            if self.magic == 0:
                # https://issues.apache.org/jira/browse/KAFKA-3160
                raise UnsupportedCodecError(
                    "LZ4 is not supported for broker version 0.8/0.9"
                )
            else:
                raw_bytes = lz4_decode(self.value)
        elif codec == self.CODEC_ZSTD:
            assert has_zstd(), "ZSTD decompression unsupported"
            raw_bytes = zstd_decode(self.value)
        else:
            raise AssertionError("This should be impossible")

        return MessageSet.decode(raw_bytes, bytes_to_read=len(raw_bytes))

    def __hash__(self) -> int:
        return hash(self.encode(recalc_crc=False))


class PartialMessage(bytes):
    def __repr__(self) -> str:
        return f"PartialMessage({self!r})"


class MessageSet:
    ITEM = Schema(("offset", Int64), ("message", Bytes))
    HEADER_SIZE = 12  # offset + message_size

    @classmethod
    def encode(
        cls,
        items: Union[io.BytesIO, Iterable[tuple[int, bytes]]],
        prepend_size: bool = True,
    ) -> bytes:
        # RecordAccumulator encodes messagesets internally
        if isinstance(items, io.BytesIO):
            size = Int32.decode(items)
            if prepend_size:
                # rewind and return all the bytes
                items.seek(items.tell() - 4)
                size += 4
            return items.read(size)

        encoded_values: list[bytes] = []
        for offset, message in items:
            encoded_values.append(Int64.encode(offset))
            encoded_values.append(Bytes.encode(message))
        encoded = b"".join(encoded_values)
        if prepend_size:
            return Bytes.encode(encoded)
        else:
            return encoded

    @classmethod
    def decode(
        cls, data: Union[io.BytesIO, bytes], bytes_to_read: Optional[int] = None
    ) -> list[Union[tuple[int, int, Message], tuple[None, None, PartialMessage]]]:
        """Compressed messages should pass in bytes_to_read (via message size)
        otherwise, we decode from data as Int32
        """
        if isinstance(data, bytes):
            data = io.BytesIO(data)
        if bytes_to_read is None:
            bytes_to_read = Int32.decode(data)

        # if FetchRequest max_bytes is smaller than the available message set
        # the server returns partial data for the final message
        # So create an internal buffer to avoid over-reading
        raw = io.BytesIO(data.read(bytes_to_read))

        items: list[
            Union[tuple[int, int, Message], tuple[None, None, PartialMessage]]
        ] = []
        try:
            while bytes_to_read:
                offset = Int64.decode(raw)
                msg_bytes = Bytes.decode(raw)
                assert msg_bytes is not None
                bytes_to_read -= 8 + 4 + len(msg_bytes)
                items.append(
                    (offset, len(msg_bytes), Message.decode(msg_bytes)),
                )
        except ValueError:
            # PartialMessage to signal that max_bytes may be too small
            items.append(
                (None, None, PartialMessage()),
            )
        return items

    @classmethod
    def repr(
        cls,
        messages: Union[
            io.BytesIO,
            list[Union[tuple[int, int, Message], tuple[None, None, PartialMessage]]],
        ],
    ) -> str:
        if isinstance(messages, io.BytesIO):
            offset = messages.tell()
            decoded = cls.decode(messages)
            messages.seek(offset)
            decoded_messages = decoded
        else:
            decoded_messages = messages
        return str([cls.ITEM.repr(m) for m in decoded_messages])
