from __future__ import annotations

import json
import struct
import typing

from ..base import BaseCodec, BaseFileIterable


class KafkaIterable(BaseFileIterable):
    """
    Apache Kafka message format reader/writer.
    Kafka message format: [offset][timestamp][key length][key][value length][value][headers]
    Simplified version for iterable reading/writing.
    """

    datamode = "binary"

    def __init__(
        self,
        filename: str = None,
        stream: typing.IO = None,
        codec: BaseCodec = None,
        mode: str = "r",
        key_name: str = "key",
        value_name: str = "value",
        include_metadata: bool = True,
        options: dict = None,
    ):
        """
        Initialize Kafka iterable.

        Args:
            key_name: Key name for the message key when reading (default: 'key')
            value_name: Key name for the message value when reading (default: 'value')
            include_metadata: Include offset, timestamp, partition in output (default: True)
        """
        if options is None:
            options = {}
        super().__init__(filename, stream, codec=codec, binary=True, mode=mode, options=options)
        self.key_name = key_name
        self.value_name = value_name
        self.include_metadata = include_metadata
        if "key_name" in options:
            self.key_name = options["key_name"]
        if "value_name" in options:
            self.value_name = options["value_name"]
        if "include_metadata" in options:
            self.include_metadata = options["include_metadata"]

        self.reset()
        pass

    def reset(self):
        """Reset iterable"""
        super().reset()
        self.pos = 0
        self.offset = 0
        if self.mode == "r":
            # File is already opened by parent class
            pass

    @staticmethod
    def id() -> str:
        return "kafka"

    @staticmethod
    def is_flatonly() -> bool:
        return False

    def read(self, skip_empty: bool = True) -> dict:
        """Read single Kafka message"""
        try:
            # Read offset (8 bytes, int64, big-endian)
            offset_bytes = self.fobj.read(8)
            if len(offset_bytes) < 8:
                raise StopIteration
            offset = struct.unpack(">q", offset_bytes)[0]

            # Read timestamp (8 bytes, int64, big-endian)
            timestamp_bytes = self.fobj.read(8)
            if len(timestamp_bytes) < 8:
                raise StopIteration
            timestamp = struct.unpack(">q", timestamp_bytes)[0]

            # Read key length (4 bytes, int32, big-endian, -1 means null)
            key_len_bytes = self.fobj.read(4)
            if len(key_len_bytes) < 4:
                raise StopIteration
            key_len = struct.unpack(">i", key_len_bytes)[0]

            # Read key
            key_data = None
            if key_len >= 0:
                key_data = self.fobj.read(key_len)
                if len(key_data) < key_len:
                    raise StopIteration

            # Read value length (4 bytes, int32, big-endian, -1 means null)
            value_len_bytes = self.fobj.read(4)
            if len(value_len_bytes) < 4:
                raise StopIteration
            value_len = struct.unpack(">i", value_len_bytes)[0]

            # Read value
            value_data = None
            if value_len >= 0:
                value_data = self.fobj.read(value_len)
                if len(value_data) < value_len:
                    raise StopIteration

            # Read headers count (2 bytes, int16, big-endian)
            headers_count_bytes = self.fobj.read(2)
            if len(headers_count_bytes) < 2:
                raise StopIteration
            headers_count = struct.unpack(">h", headers_count_bytes)[0]

            # Read headers
            headers = {}
            for _ in range(headers_count):
                # Read header key length (2 bytes, int16, big-endian)
                hkey_len_bytes = self.fobj.read(2)
                if len(hkey_len_bytes) < 2:
                    raise StopIteration
                hkey_len = struct.unpack(">h", hkey_len_bytes)[0]

                # Read header key
                hkey = self.fobj.read(hkey_len).decode("utf-8")

                # Read header value length (2 bytes, int16, big-endian)
                hval_len_bytes = self.fobj.read(2)
                if len(hval_len_bytes) < 2:
                    raise StopIteration
                hval_len = struct.unpack(">h", hval_len_bytes)[0]

                # Read header value
                hval = self.fobj.read(hval_len)
                headers[hkey] = hval.decode("utf-8", errors="replace")

            # Parse key and value
            key_obj = None
            if key_data:
                try:
                    key_str = key_data.decode("utf-8")
                    try:
                        key_obj = json.loads(key_str)
                    except json.JSONDecodeError:
                        key_obj = key_str
                except UnicodeDecodeError:
                    import base64

                    key_obj = base64.b64encode(key_data).decode("utf-8")

            value_obj = None
            if value_data:
                try:
                    value_str = value_data.decode("utf-8")
                    try:
                        value_obj = json.loads(value_str)
                    except json.JSONDecodeError:
                        value_obj = value_str
                except UnicodeDecodeError:
                    import base64

                    value_obj = base64.b64encode(value_data).decode("utf-8")

            result = {}
            if self.include_metadata:
                result["offset"] = offset
                result["timestamp"] = timestamp
                if headers:
                    result["headers"] = headers

            if key_obj is not None:
                result[self.key_name] = key_obj
            if value_obj is not None:
                result[self.value_name] = value_obj

            self.pos += 1
            self.offset = offset
            return result

        except StopIteration:
            raise StopIteration from None
        except Exception as e:
            raise ValueError(f"Error reading Kafka message: {e}") from e

    def read_bulk(self, num: int = 10) -> list[dict]:
        """Read bulk Kafka messages"""
        chunk = []
        for _n in range(0, num):
            try:
                chunk.append(self.read())
            except StopIteration:
                break
        return chunk

    def write(self, record: dict):
        """Write single Kafka message"""
        # Extract key and value from record
        key_obj = record.get(self.key_name)
        value_obj = record.get(self.value_name)

        # If not found, try to infer from record structure
        if key_obj is None and value_obj is None:
            if len(record) == 2 and self.include_metadata:
                # Skip metadata fields
                items = [(k, v) for k, v in record.items() if k not in ["offset", "timestamp", "headers"]]
                if len(items) == 2:
                    key_obj = items[0][1]
                    value_obj = items[1][1]
            elif len(record) == 1:
                value_obj = list(record.values())[0]

        # Get metadata
        offset = record.get("offset", self.offset + 1)
        timestamp = record.get("timestamp", 0)
        headers = record.get("headers", {})

        # Convert to bytes
        key_data = None
        if key_obj is not None:
            if isinstance(key_obj, str):
                key_data = key_obj.encode("utf-8")
            elif isinstance(key_obj, (dict, list)):
                key_data = json.dumps(key_obj, ensure_ascii=False).encode("utf-8")
            else:
                key_data = str(key_obj).encode("utf-8")

        if isinstance(value_obj, str):
            value_data = value_obj.encode("utf-8")
        elif isinstance(value_obj, (dict, list)):
            value_data = json.dumps(value_obj, ensure_ascii=False).encode("utf-8")
        elif value_obj is not None:
            value_data = str(value_obj).encode("utf-8")
        else:
            value_data = None

        # Write offset (8 bytes, int64, big-endian)
        self.fobj.write(struct.pack(">q", offset))

        # Write timestamp (8 bytes, int64, big-endian)
        self.fobj.write(struct.pack(">q", timestamp))

        # Write key length (4 bytes, int32, big-endian, -1 for null)
        if key_data:
            self.fobj.write(struct.pack(">i", len(key_data)))
            self.fobj.write(key_data)
        else:
            self.fobj.write(struct.pack(">i", -1))

        # Write value length (4 bytes, int32, big-endian, -1 for null)
        if value_data:
            self.fobj.write(struct.pack(">i", len(value_data)))
            self.fobj.write(value_data)
        else:
            self.fobj.write(struct.pack(">i", -1))

        # Write headers count (2 bytes, int16, big-endian)
        self.fobj.write(struct.pack(">h", len(headers)))

        # Write headers
        for hkey, hval in headers.items():
            hkey_bytes = hkey.encode("utf-8")
            hval_bytes = str(hval).encode("utf-8")
            self.fobj.write(struct.pack(">h", len(hkey_bytes)))
            self.fobj.write(hkey_bytes)
            self.fobj.write(struct.pack(">h", len(hval_bytes)))
            self.fobj.write(hval_bytes)

        self.offset = offset

    def write_bulk(self, records: list[dict]):
        """Write bulk Kafka messages"""
        for record in records:
            self.write(record)
