from __future__ import annotations
import typing
import struct
import json
from ..base import BaseFileIterable, BaseCodec


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={}):
        """
        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)
        """
        super(KafkaIterable, self).__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(KafkaIterable, self).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
        except Exception as e:
            raise ValueError(f"Error reading Kafka message: {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)
