"""
Utility functions for converting between Python objects and Protobuf messages.

This module provides helper functions to:
- Convert internal `juturna.components.Message` and its Payloads
  to their corresponding Protobuf representations in `protos/payloads.proto`.
- Deserialize `ProtoMessage` and `ProtoPayload` back into Python objects.
- Create/deserialize `ProtoEnvelope` for wrapping messages in the wire protocol.

These utilities are essential for the serialization layer of the Remotizer.
"""

import numpy as np
import base64
import uuid
import time
from typing import Any
import logging
from datetime import datetime, date
from decimal import Decimal


from juturna.components import Message
from juturna.payloads import (
    AudioPayload,
    ImagePayload,
    VideoPayload,
    BytesPayload,
    ObjectPayload,
    Batch,
)

from juturna.remotizer.c_protos.payloads_pb2 import (
    ProtoMessage,
    ProtoEnvelope,
    AudioProtoPayload,
    ImageProtoPayload,
    VideoProtoPayload,
    BytesProtoPayload,
    ObjectProtoPayload,
    BatchProto,
)
from google.protobuf.struct_pb2 import Struct
from google.protobuf.json_format import MessageToDict

logger = logging.getLogger('jt.remotizer.utils')

# ============================================================================
# PYTHON → PROTOBUF CONVERTERS
# ============================================================================


def _audio_to_proto(audio: AudioPayload) -> AudioProtoPayload:
    """Convert Python AudioPayload to Protobuf AudioProtoPayload"""
    proto = AudioProtoPayload()

    # Convert numpy array to bytes
    proto.audio_data = audio.audio.tobytes()
    proto.dtype = str(audio.audio.dtype)
    proto.shape.extend(audio.audio.shape)

    # Copy metadata
    proto.sampling_rate = audio.sampling_rate
    proto.channels = audio.channels
    proto.start = audio.start
    proto.end = audio.end

    return proto


def _image_to_proto(image: ImagePayload) -> ImageProtoPayload:
    """Convert Python ImagePayload to Protobuf ImageProtoPayload"""
    proto = ImageProtoPayload()

    # Convert numpy array to bytes
    proto.image_data = image.image.tobytes()
    proto.dtype = str(image.image.dtype)

    # Copy metadata
    proto.width = image.width
    proto.height = image.height
    proto.depth = image.depth
    proto.pixel_format = image.pixel_format
    proto.timestamp = image.timestamp

    return proto


def _video_to_proto(video: VideoPayload) -> VideoProtoPayload:
    """Convert Python VideoPayload to Protobuf VideoProtoPayload"""
    proto = VideoProtoPayload()

    # Convert each frame
    for frame in video.video:
        frame_proto = _image_to_proto(frame)
        proto.frames.append(frame_proto)

    # Copy metadata
    proto.frames_per_second = video.frames_per_second
    proto.start = video.start
    proto.codec = video.codec
    proto.end = video.end

    return proto


def _bytes_to_proto(bytes_payload: BytesPayload) -> BytesProtoPayload:
    """Convert Python BytesPayload to Protobuf BytesProtoPayload"""
    proto = BytesProtoPayload()
    proto.cnt = bytes_payload.cnt
    proto.size = len(bytes_payload.cnt)
    return proto


def _object_to_proto(obj: ObjectPayload) -> ObjectProtoPayload:
    """Convert Python ObjectPayload (dict) to Protobuf ObjectProtoPayload"""
    proto = ObjectProtoPayload()

    # Convert dict to Struct
    struct = Struct()
    struct.update(dict(obj))
    proto.data.CopyFrom(struct)

    return proto


def _batch_to_proto(batch: Batch) -> BatchProto:
    """Convert Python Batch to Protobuf BatchProto"""
    proto = BatchProto()
    for message in batch.messages:
        message_proto = message_to_proto(message)
        proto.messages.append(message_proto)
    return proto


PROTOBUF_PAYLOAD_TYPE_MAP = {
    AudioPayload: _audio_to_proto,
    ImagePayload: _image_to_proto,
    VideoPayload: _video_to_proto,
    BytesPayload: _bytes_to_proto,
    ObjectPayload: _object_to_proto,
    Batch: _batch_to_proto,
}


def message_to_proto(message: Message) -> ProtoMessage:
    """
    Convert Python Message to Protobuf ProtoMessage

    Handles all payload types:
    - AudioPayload → AudioProtoPayload
    - ImagePayload → ImageProtoPayload
    - VideoPayload → VideoProtoPayload
    - BytesPayload → BytesProtoPayload
    - ObjectPayload → ObjectProtoPayload
    - Batch → BatchProto

    Parameters
    ----------
    message : Message
        Python Message object with any payload type

    Returns
    -------
    ProtoMessage
        Protobuf message ready for serialization

    """
    proto = ProtoMessage()

    proto.created_at = message.created_at
    proto.creator = message.creator
    proto.version = message.version
    proto.id = message.id

    proto.meta.update(sanitize_struct_for_proto(message.meta))
    proto.timers.update(dict(message.timers))

    if message.payload is not None:
        protocol_converter = PROTOBUF_PAYLOAD_TYPE_MAP.get(
            type(message.payload)
        )
        payload_proto = protocol_converter(message.payload)
        proto.payload.Pack(payload_proto)
    return proto


# ============================================================================
# PROTOBUF →  PYTHON CONVERTERS
# ============================================================================


def deserialize_message(message: ProtoMessage) -> Message:
    """
    Convert Protobuf ProtoMessage to Python Message

    Automatically unpacks and converts the payload based on its type:
    - AudioProtoPayload → AudioPayload
    - ImageProtoPayload → ImagePayload
    - VideoProtoPayload → VideoPayload
    - BytesProtoPayload → BytesPayload
    - ObjectProtoPayload → ObjectPayload
    - BatchProto → Batch

    Parameters
    ----------
    message : ProtoMessage
        Protobuf message to convert

    Returns
    -------
    Message
        Python Message object with appropriate payload type

    """
    message_obj = Message(
        creator=message.creator,
        version=message.version,
        payload=None,  # to be filled below
    )
    message_obj.created_at = message.created_at
    message_obj.meta.update(dict(message.meta))
    message_obj.timers.update(dict(message.timers))
    message_obj.id = message.id

    if message.payload.Is(AudioProtoPayload.DESCRIPTOR):
        audio = AudioProtoPayload()
        message.payload.Unpack(audio)
        message_obj.payload = _deserialize_audio_payload(audio)

    elif message.payload.Is(ImageProtoPayload.DESCRIPTOR):
        image = ImageProtoPayload()
        message.payload.Unpack(image)
        message_obj.payload = _deserialize_image_payload(image)

    elif message.payload.Is(VideoProtoPayload.DESCRIPTOR):
        video = VideoProtoPayload()
        message.payload.Unpack(video)
        message_obj.payload = _deserialize_video_payload(video)

    elif message.payload.Is(BytesProtoPayload.DESCRIPTOR):
        bytes_payload = BytesProtoPayload()
        message.payload.Unpack(bytes_payload)
        message_obj.payload = _deserialize_bytes_payload(bytes_payload)

    elif message.payload.Is(ObjectProtoPayload.DESCRIPTOR):
        obj = ObjectProtoPayload()
        message.payload.Unpack(obj)
        message_obj.payload = _deserialize_object_payload(obj)

    elif message.payload.Is(BatchProto.DESCRIPTOR):
        batch = BatchProto()
        message.payload.Unpack(batch)
        message_obj.payload = _deserialize_batch_payload(batch)

    return message_obj


def _deserialize_audio_payload(payload: AudioProtoPayload) -> AudioPayload:
    """Deserialize AudioProtoPayload to AudioPayload with numpy array"""
    audio_data = np.frombuffer(payload.audio_data, dtype=payload.dtype)
    audio_data = audio_data.reshape(payload.shape)

    return AudioPayload(
        audio=audio_data,
        sampling_rate=payload.sampling_rate,
        channels=payload.channels,
        start=payload.start,
        end=payload.end,
    )


def _deserialize_image_payload(payload: ImageProtoPayload) -> ImagePayload:
    """Deserialize ImageProtoPayload to ImagePayload with numpy array"""
    image_data = np.frombuffer(payload.image_data, dtype=payload.dtype)

    if payload.depth == 1:
        shape = (payload.height, payload.width)
    else:
        shape = (payload.height, payload.width, payload.depth)

    image_data = image_data.reshape(shape)

    return ImagePayload(
        image=image_data,
        width=payload.width,
        height=payload.height,
        depth=payload.depth,
        pixel_format=payload.pixel_format,
        timestamp=payload.timestamp,
    )


def _deserialize_video_payload(payload: VideoProtoPayload) -> VideoPayload:
    """Deserialize VideoProtoPayload to VideoPayload with numpy arrays list"""
    frames = [_deserialize_image_payload(frame) for frame in payload.frames]

    return VideoPayload(
        video=frames,
        frames_per_second=payload.frames_per_second,
        codec=payload.codec,
        start=payload.start,
        end=payload.end,
    )


def _deserialize_bytes_payload(payload: BytesProtoPayload) -> BytesPayload:
    """Deserialize BytesProtoPayload to BytesPayload"""
    return BytesPayload(
        cnt=payload.cnt,
    )


def _deserialize_object_payload(payload: ObjectProtoPayload) -> ObjectPayload:
    """Deserialize ObjectProtoPayload (Struct) to ObjectPayload (dict)"""
    proto_dict = MessageToDict(payload)
    return ObjectPayload.from_dict(proto_dict.get('data', {}))


def _deserialize_batch_payload(payload: BatchProto) -> Batch:
    """Deserialize BatchProto to Batch"""
    return Batch(messages=tuple(payload.messages))


# ! FIXME: probably we should change the request_type and response_type
# ! signature: they are meaningful only in the context of requests,
# ! for responses they should be empty, or with a different naming
def create_envelope(
    message: ProtoMessage,
    configuration: dict[str, Any],
    metadata: dict[str, Any],
    creator: str,
    id: str = str(uuid.uuid4()),
    priority: int = 1,
    timeout: int = 30,
    response_to: str = None,
    response_type: str = '',
    request_type: str = '',
) -> ProtoEnvelope:
    """Create envelope around message"""
    assert message is not None

    envelope = ProtoEnvelope()
    envelope.id = id
    envelope.sender = creator
    envelope.created_at = time.time()
    envelope.ttl = int(timeout)
    envelope.request_type = request_type
    envelope.response_type = response_type
    envelope.priority = priority
    envelope.response_to = response_to if response_to is not None else ''
    envelope.configuration.update(configuration)
    envelope.metadata.update(metadata)
    envelope.message.CopyFrom(message)

    return envelope


def deserialize_envelope(envelope: ProtoEnvelope) -> dict[str, Any]:
    """Deserialize ProtoEnvelope to dictionary"""
    message = deserialize_message(envelope.message)
    envelope_dict = {
        'id': envelope.id,
        'sender': envelope.sender,
        'response_to': envelope.response_to,
        'ttl': envelope.ttl,
        'request_type': envelope.request_type,
        'response_type': envelope.response_type,
        'message': message,
    }
    return envelope_dict


def to_primitive(obj: Any) -> Any:
    """
    Convert any object to JSON-compatible primitives.
    Returns None for non-convertible objects.
    """
    if obj is None or isinstance(obj, (bool, int, float, str)):
        return obj

    if isinstance(obj, (np.integer, np.floating)):
        return obj.item()
    if isinstance(obj, np.ndarray):
        return obj.tolist()

    if isinstance(obj, (datetime, date)):
        return obj.isoformat()

    if isinstance(obj, Decimal):
        return float(obj)

    if isinstance(obj, bytes):
        try:
            return base64.b64encode(obj).decode('ascii')
        except UnicodeDecodeError:
            return None

    if isinstance(obj, dict):
        result = {}
        for k, v in obj.items():
            key = str(k) if not isinstance(k, str) else k
            value = to_primitive(v)
            if value is not None or v is None:
                result[key] = value
        return result

    if isinstance(obj, (list, tuple)):
        result = []
        for item in obj:
            value = to_primitive(item)
            if value is not None or item is None:
                result.append(value)
        return result

    # Try __dict__ for custom objects
    if hasattr(obj, '__dict__'):
        return to_primitive(obj.__dict__)

    logger.warning(f'dropping non-serializable object: {type(obj).__name__}')
    return None


def sanitize_struct_for_proto(meta: dict[str, Any]) -> dict[str, Any]:
    """
    Sanitize a metadata dictionary for protobuf Struct serialization.

    Args:
        meta: Original metadata dictionary

    Returns:
        Sanitized dictionary compatible with google.protobuf.Struct

    """
    if not meta:
        return {}

    logger.info(f'sanitizing metadata with {len(meta)} entries')
    sanitized = to_primitive(dict(meta))

    if sanitized is None:
        logger.warning(
            'entire metadata is non-serializable, returning empty dict'
        )
        return {}

    logger.info(f'sanitized metadata: {sanitized}')

    return sanitized
