"""Message data structure."""

import asyncio
import json
from dataclasses import dataclass, field
from io import BytesIO
from typing import Dict, Any, Optional, List, TYPE_CHECKING

from .conversation import Conversation
from .conversation_client import ConversationClient

if TYPE_CHECKING:
    from .memory import Memory
    from .memory_v2 import MemoryV2
    from .memory_client import MemoryClient
    from .memory_v2_client import MemoryV2Client
    from .attachment_client import AttachmentClient, AttachmentInfo


@dataclass
class Message:
    """Represents an incoming message."""

    content: str
    conversation_id: str
    user_id: str
    headers: Dict[str, str]
    raw_body: Dict[str, Any]
    conversation: Optional[Conversation] = None
    memory: Optional["Memory"] = None
    memory_v2: Optional["MemoryV2"] = None
    metadata: Optional[Dict[str, Any]] = None
    schema: Optional[Dict[str, Any]] = None
    # Keep singular for backward compatibility
    attachment: Optional["AttachmentInfo"] = None
    # Array of all attachments
    attachments: List["AttachmentInfo"] = field(default_factory=list)
    _attachment_client: Optional["AttachmentClient"] = field(default=None, repr=False)

    def __post_init__(self):
        """Ensure backward compatibility between attachment and attachments."""
        # If attachments is empty but attachment exists, populate attachments
        if not self.attachments and self.attachment:
            self.attachments = [self.attachment]
        # If attachment is None but attachments exist, set attachment to first
        if self.attachment is None and self.attachments:
            self.attachment = self.attachments[0]

    @classmethod
    def from_kafka_message(
        cls,
        body: Dict[str, Any],
        headers: Dict[str, str],
        conversation_client: Optional[ConversationClient] = None,
        memory_client: Optional["MemoryClient"] = None,
        memory_v2_client: Optional["MemoryV2Client"] = None,
        attachment_client: Optional["AttachmentClient"] = None,
        agent_name: Optional[str] = None,
    ):
        """Create Message from Kafka message."""
        from .memory import Memory
        from .attachment_client import AttachmentInfo

        conversation_id = headers.get("conversationId")
        user_id = headers.get("userId", "anonymous")

        # Parse metadata from header
        metadata = None
        metadata_header = headers.get("messageMetadata")
        if metadata_header:
            try:
                metadata = json.loads(metadata_header)
            except json.JSONDecodeError:
                pass

        # Parse schema from header
        schema = None
        schema_header = headers.get("messageSchema")
        if schema_header:
            try:
                schema = json.loads(schema_header)
            except json.JSONDecodeError:
                pass

        # Parse attachments from header or body (prefer plural, fallback to singular)
        attachments: List[AttachmentInfo] = []

        # Try parsing plural 'attachments' header first (new format)
        attachments_header = headers.get("attachments")
        if attachments_header:
            try:
                attachments_data = json.loads(attachments_header)
                attachments = [AttachmentInfo.from_dict(a) for a in attachments_data]
            except json.JSONDecodeError:
                pass

        # Fallback to body attachments (plural)
        if not attachments and body.get("attachments"):
            try:
                attachments = [
                    AttachmentInfo.from_dict(a) for a in body["attachments"]
                ]
            except (TypeError, KeyError):
                pass

        # BACKWARD COMPATIBILITY: Fallback to singular 'attachment' header
        if not attachments:
            attachment_header = headers.get("attachment")
            if attachment_header:
                try:
                    attachment_data = json.loads(attachment_header)
                    attachments = [AttachmentInfo.from_dict(attachment_data)]
                except json.JSONDecodeError:
                    pass

        # Fallback to singular body attachment
        if not attachments and body.get("attachment"):
            try:
                attachments = [AttachmentInfo.from_dict(body["attachment"])]
            except (TypeError, KeyError):
                pass

        conversation = None
        if conversation_id and conversation_client:
            conversation = Conversation(conversation_id, conversation_client, agent_name)

        memory = None
        if memory_client and user_id and conversation_id:
            memory = Memory(user_id, conversation_id, memory_client)

        memory_v2 = None
        if memory_v2_client and user_id and conversation_id:
            from .memory_v2 import MemoryV2
            memory_v2 = MemoryV2(user_id, conversation_id, memory_v2_client)

        return cls(
            content=body.get("content", ""),
            conversation_id=conversation_id,
            user_id=user_id,
            headers=headers,
            raw_body=body,
            conversation=conversation,
            memory=memory,
            memory_v2=memory_v2,
            metadata=metadata,
            schema=schema,
            attachment=attachments[0] if attachments else None,
            attachments=attachments,
            _attachment_client=attachment_client,
        )

    def has_attachment(self) -> bool:
        """Check if message has at least one attachment."""
        return len(self.attachments) > 0

    def has_attachments(self) -> bool:
        """Check if message has any attachments. Alias for has_attachment()."""
        return len(self.attachments) > 0

    def get_attachment_count(self) -> int:
        """Get the number of attachments."""
        return len(self.attachments)

    def get_attachments(self) -> List["AttachmentInfo"]:
        """Get all attachments."""
        from .attachment_client import AttachmentInfo

        return list(self.attachments)

    async def get_attachment_bytes(self) -> bytes:
        """
        Download first attachment into memory as bytes.
        For backward compatibility with single-attachment code.

        Returns:
            File content as bytes

        Raises:
            ValueError: If no attachment is available
        """
        if not self.attachments:
            raise ValueError("No attachment available")
        if not self._attachment_client:
            raise ValueError("AttachmentClient not available")
        return await self._attachment_client.download(self.attachments[0].url)

    async def get_attachment_base64(self) -> str:
        """
        Get first attachment as base64 string (for LLM vision APIs).
        For backward compatibility with single-attachment code.

        Returns:
            Base64 encoded string of the file content

        Raises:
            ValueError: If no attachment is available
        """
        if not self.attachments:
            raise ValueError("No attachment available")
        if not self._attachment_client:
            raise ValueError("AttachmentClient not available")
        return await self._attachment_client.download_as_base64(self.attachments[0].url)

    async def get_attachment_buffer(self) -> BytesIO:
        """
        Get first attachment as BytesIO buffer (for document processing libraries).
        For backward compatibility with single-attachment code.

        Returns:
            BytesIO buffer containing the file content

        Raises:
            ValueError: If no attachment is available
        """
        if not self.attachments:
            raise ValueError("No attachment available")
        if not self._attachment_client:
            raise ValueError("AttachmentClient not available")
        return await self._attachment_client.download_to_buffer(self.attachments[0].url)

    async def get_attachment_bytes_at(self, index: int) -> bytes:
        """
        Download a specific attachment by index as bytes.

        Args:
            index: Zero-based index of the attachment

        Returns:
            File content as bytes

        Raises:
            ValueError: If index is out of bounds or no attachment client
        """
        if index < 0 or index >= len(self.attachments):
            raise ValueError(
                f"Invalid attachment index: {index}. "
                f"Valid range: 0-{len(self.attachments) - 1}"
            )
        if not self._attachment_client:
            raise ValueError("AttachmentClient not available")
        return await self._attachment_client.download(self.attachments[index].url)

    async def get_attachment_base64_at(self, index: int) -> str:
        """
        Get a specific attachment by index as base64 string.

        Args:
            index: Zero-based index of the attachment

        Returns:
            Base64 encoded string of the file content

        Raises:
            ValueError: If index is out of bounds or no attachment client
        """
        if index < 0 or index >= len(self.attachments):
            raise ValueError(
                f"Invalid attachment index: {index}. "
                f"Valid range: 0-{len(self.attachments) - 1}"
            )
        if not self._attachment_client:
            raise ValueError("AttachmentClient not available")
        return await self._attachment_client.download_as_base64(
            self.attachments[index].url
        )

    async def get_attachment_buffer_at(self, index: int) -> BytesIO:
        """
        Get a specific attachment by index as BytesIO buffer.

        Args:
            index: Zero-based index of the attachment

        Returns:
            BytesIO buffer containing the file content

        Raises:
            ValueError: If index is out of bounds or no attachment client
        """
        if index < 0 or index >= len(self.attachments):
            raise ValueError(
                f"Invalid attachment index: {index}. "
                f"Valid range: 0-{len(self.attachments) - 1}"
            )
        if not self._attachment_client:
            raise ValueError("AttachmentClient not available")
        return await self._attachment_client.download_to_buffer(
            self.attachments[index].url
        )

    async def get_all_attachment_bytes(self) -> List[bytes]:
        """
        Download all attachments as bytes.

        Returns:
            List of file contents as bytes in the same order as attachments

        Raises:
            ValueError: If no attachment client available
        """
        if not self._attachment_client:
            raise ValueError("AttachmentClient not available")
        return await asyncio.gather(
            *[self._attachment_client.download(a.url) for a in self.attachments]
        )

    async def get_all_attachment_base64(self) -> List[str]:
        """
        Get all attachments as base64 strings.

        Returns:
            List of base64 encoded strings in the same order as attachments

        Raises:
            ValueError: If no attachment client available
        """
        if not self._attachment_client:
            raise ValueError("AttachmentClient not available")
        return await asyncio.gather(
            *[
                self._attachment_client.download_as_base64(a.url)
                for a in self.attachments
            ]
        )

    async def get_all_attachment_buffers(self) -> List[BytesIO]:
        """
        Get all attachments as BytesIO buffers.

        Returns:
            List of BytesIO buffers in the same order as attachments

        Raises:
            ValueError: If no attachment client available
        """
        if not self._attachment_client:
            raise ValueError("AttachmentClient not available")
        return await asyncio.gather(
            *[
                self._attachment_client.download_to_buffer(a.url)
                for a in self.attachments
            ]
        )
