r2rpy/py/core/providers/database/conversations.py

import csv
import json
import logging
import tempfile
from datetime import datetime
from typing import IO, Any, Optional
from uuid import UUID, uuid4

from fastapi import HTTPException

from core.base import Handler, Message, R2RException
from shared.api.models.management.responses import (
    ConversationResponse,
    MessageResponse,
)

from .base import PostgresConnectionManager

logger = logging.getLogger(__name__)


def _validate_image_size(
    message: Message, max_size_bytes: int = 5 * 1024 * 1024
) -> None:
    """
    Validates that images in a message don't exceed the maximum allowed size.

    Args:
        message: Message object to validate
        max_size_bytes: Maximum allowed size for base64-encoded images (default: 5MB)

    Raises:
        R2RException: If image is too large
    """
    if (
        hasattr(message, "image_data")
        and message.image_data
        and "data" in message.image_data
    ):
        base64_data = message.image_data["data"]

        # Calculate approximate decoded size (base64 increases size by ~33%)
        # The formula is: decoded_size = encoded_size * 3/4
        estimated_size_bytes = len(base64_data) * 0.75

        if estimated_size_bytes > max_size_bytes:
            raise R2RException(
                status_code=413,  # Payload Too Large
                message=f"Image too large: {estimated_size_bytes / 1024 / 1024:.2f}MB exceeds the maximum allowed size of {max_size_bytes / 1024 / 1024:.2f}MB",
            )


def _json_default(obj: Any) -> str:
    """Default handler for objects not serializable by the standard json
    encoder."""
    if isinstance(obj, datetime):
        # Return ISO8601 string
        return obj.isoformat()
    elif isinstance(obj, UUID):
        # Convert UUID to string
        return str(obj)
    # If you have other special types, handle them here...
    # e.g. decimal.Decimal -> str(obj)

    # If we get here, raise an error or just default to string:
    raise TypeError(f"Type {type(obj)} not serializable")


def safe_dumps(obj: Any) -> str:
    """Wrap `json.dumps` with a default that serializes UUID and datetime."""
    return json.dumps(obj, default=_json_default)


class PostgresConversationsHandler(Handler):
    def __init__(
        self, project_name: str, connection_manager: PostgresConnectionManager
    ):
        self.project_name = project_name
        self.connection_manager = connection_manager

    async def create_tables(self):
        create_conversations_query = f"""
        CREATE TABLE IF NOT EXISTS {self._get_table_name("conversations")} (
            id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
            user_id UUID,
            created_at TIMESTAMPTZ DEFAULT NOW(),
            name TEXT
        );
        """

        create_messages_query = f"""
        CREATE TABLE IF NOT EXISTS {self._get_table_name("messages")} (
            id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
            conversation_id UUID NOT NULL,
            parent_id UUID,
            content JSONB,
            metadata JSONB,
            created_at TIMESTAMPTZ DEFAULT NOW(),
            FOREIGN KEY (conversation_id) REFERENCES {self._get_table_name("conversations")}(id),
            FOREIGN KEY (parent_id) REFERENCES {self._get_table_name("messages")}(id)
        );
        """
        await self.connection_manager.execute_query(create_conversations_query)
        await self.connection_manager.execute_query(create_messages_query)

    async def create_conversation(
        self,
        user_id: Optional[UUID] = None,
        name: Optional[str] = None,
    ) -> ConversationResponse:
        query = f"""
            INSERT INTO {self._get_table_name("conversations")} (user_id, name)
            VALUES ($1, $2)
            RETURNING id, extract(epoch from created_at) as created_at_epoch
        """
        try:
            result = await self.connection_manager.fetchrow_query(
                query, [user_id, name]
            )

            return ConversationResponse(
                id=result["id"],
                created_at=result["created_at_epoch"],
                user_id=user_id or None,
                name=name or None,
            )
        except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"Failed to create conversation: {str(e)}",
            ) from e

    async def get_conversations_overview(
        self,
        offset: int,
        limit: int,
        filter_user_ids: Optional[list[UUID]] = None,
        conversation_ids: Optional[list[UUID]] = None,
    ) -> dict[str, Any]:
        conditions = []
        params: list = []
        param_index = 1

        if filter_user_ids:
            conditions.append(f"""
                c.user_id IN (
                    SELECT id
                    FROM {self.project_name}.users
                    WHERE id = ANY(${param_index})
                )
            """)
            params.append(filter_user_ids)
            param_index += 1

        if conversation_ids:
            conditions.append(f"c.id = ANY(${param_index})")
            params.append(conversation_ids)
            param_index += 1

        where_clause = (
            "WHERE " + " AND ".join(conditions) if conditions else ""
        )

        query = f"""
            WITH conversation_overview AS (
                SELECT c.id,
                    extract(epoch from c.created_at) as created_at_epoch,
                    c.user_id,
                    c.name
                FROM {self._get_table_name("conversations")} c
                {where_clause}
            ),
            counted_overview AS (
                SELECT *,
                    COUNT(*) OVER() AS total_entries
                FROM conversation_overview
            )
            SELECT * FROM counted_overview
            ORDER BY created_at_epoch DESC
            OFFSET ${param_index}
        """
        params.append(offset)
        param_index += 1

        if limit != -1:
            query += f" LIMIT ${param_index}"
            params.append(limit)

        results = await self.connection_manager.fetch_query(query, params)

        if not results:
            return {"results": [], "total_entries": 0}

        total_entries = results[0]["total_entries"]
        conversations = [
            {
                "id": str(row["id"]),
                "created_at": row["created_at_epoch"],
                "user_id": str(row["user_id"]) if row["user_id"] else None,
                "name": row["name"] or None,
            }
            for row in results
        ]

        return {"results": conversations, "total_entries": total_entries}

    async def add_message(
        self,
        conversation_id: UUID,
        content: Message,
        parent_id: Optional[UUID] = None,
        metadata: Optional[dict] = None,
        max_image_size_bytes: int = 5 * 1024 * 1024,  # 5MB default
    ) -> MessageResponse:
        # Validate image size
        try:
            _validate_image_size(content, max_image_size_bytes)
        except R2RException:
            # Re-raise validation exceptions
            raise
        except Exception as e:
            # Handle unexpected errors during validation
            logger.error(f"Error validating image: {str(e)}")
            raise R2RException(
                status_code=400, message=f"Invalid image data: {str(e)}"
            ) from e

        # 1) Validate that conversation and parent exist (existing code)
        conv_check_query = f"""
            SELECT 1 FROM {self._get_table_name("conversations")}
            WHERE id = $1
        """
        conv_row = await self.connection_manager.fetchrow_query(
            conv_check_query, [conversation_id]
        )
        if not conv_row:
            raise R2RException(
                status_code=404,
                message=f"Conversation {conversation_id} not found.",
            )

        if parent_id:
            parent_check_query = f"""
                SELECT 1 FROM {self._get_table_name("messages")}
                WHERE id = $1 AND conversation_id = $2
            """
            parent_row = await self.connection_manager.fetchrow_query(
                parent_check_query, [parent_id, conversation_id]
            )
            if not parent_row:
                raise R2RException(
                    status_code=404,
                    message=f"Parent message {parent_id} not found in conversation {conversation_id}.",
                )

        # 2) Add image info to metadata for tracking/analytics if images are present
        metadata = metadata or {}
        if hasattr(content, "image_url") and content.image_url:
            metadata["has_image"] = True
            metadata["image_type"] = "url"
        elif hasattr(content, "image_data") and content.image_data:
            metadata["has_image"] = True
            metadata["image_type"] = "base64"
            # Don't store the actual base64 data in metadata as it would be redundant

        # 3) Convert the content & metadata to JSON strings
        message_id = uuid4()
        # Using safe_dumps to handle any type of serialization
        content_str = safe_dumps(content.model_dump())
        metadata_str = safe_dumps(metadata)

        # 4) Insert the message (existing code)
        query = f"""
            INSERT INTO {self._get_table_name("messages")}
            (id, conversation_id, parent_id, content, created_at, metadata)
            VALUES ($1, $2, $3, $4::jsonb, NOW(), $5::jsonb)
            RETURNING id
        """
        inserted = await self.connection_manager.fetchrow_query(
            query,
            [
                message_id,
                conversation_id,
                parent_id,
                content_str,
                metadata_str,
            ],
        )
        if not inserted:
            raise R2RException(
                status_code=500, message="Failed to insert message."
            )

        return MessageResponse(id=message_id, message=content)

    async def edit_message(
        self,
        message_id: UUID,
        new_content: str | None = None,
        additional_metadata: dict | None = None,
    ) -> dict[str, Any]:
        # Get the original message
        query = f"""
            SELECT conversation_id, parent_id, content, metadata, created_at
            FROM {self._get_table_name("messages")}
            WHERE id = $1
        """
        row = await self.connection_manager.fetchrow_query(query, [message_id])
        if not row:
            raise R2RException(
                status_code=404,
                message=f"Message {message_id} not found.",
            )

        old_content = json.loads(row["content"])
        old_metadata = json.loads(row["metadata"])

        if new_content is not None:
            old_message = Message(**old_content)
            edited_message = Message(
                role=old_message.role,
                content=new_content,
                name=old_message.name,
                function_call=old_message.function_call,
                tool_calls=old_message.tool_calls,
                # Preserve image content if it exists
                image_url=getattr(old_message, "image_url", None),
                image_data=getattr(old_message, "image_data", None),
            )
            content_to_save = edited_message.model_dump()
        else:
            content_to_save = old_content

        additional_metadata = additional_metadata or {}

        new_metadata = {
            **old_metadata,
            **additional_metadata,
            "edited": (
                True
                if new_content is not None
                else old_metadata.get("edited", False)
            ),
        }

        # Update message without changing the timestamp
        update_query = f"""
            UPDATE {self._get_table_name("messages")}
            SET content = $1::jsonb,
                metadata = $2::jsonb,
                created_at = $3
            WHERE id = $4
            RETURNING id
        """
        updated = await self.connection_manager.fetchrow_query(
            update_query,
            [
                json.dumps(content_to_save),
                json.dumps(new_metadata),
                row["created_at"],
                message_id,
            ],
        )
        if not updated:
            raise R2RException(
                status_code=500, message="Failed to update message."
            )

        return {
            "id": str(message_id),
            "message": (
                Message(**content_to_save)
                if isinstance(content_to_save, dict)
                else content_to_save
            ),
            "metadata": new_metadata,
        }

    async def update_message_metadata(
        self, message_id: UUID, metadata: dict
    ) -> None:
        # Fetch current metadata
        query = f"""
            SELECT metadata FROM {self._get_table_name("messages")}
            WHERE id = $1
        """
        row = await self.connection_manager.fetchrow_query(query, [message_id])
        if not row:
            raise R2RException(
                status_code=404, message=f"Message {message_id} not found."
            )

        current_metadata = json.loads(row["metadata"]) or {}
        updated_metadata = {**current_metadata, **metadata}

        update_query = f"""
            UPDATE {self._get_table_name("messages")}
            SET metadata = $1::jsonb
            WHERE id = $2
        """
        await self.connection_manager.execute_query(
            update_query, [json.dumps(updated_metadata), message_id]
        )

    async def get_conversation(
        self,
        conversation_id: UUID,
        filter_user_ids: Optional[list[UUID]] = None,
    ) -> list[MessageResponse]:
        # Existing validation code remains the same
        conditions = ["c.id = $1"]
        params: list = [conversation_id]

        if filter_user_ids:
            param_index = 2
            conditions.append(f"""
                c.user_id IN (
                    SELECT id
                    FROM {self.project_name}.users
                    WHERE id = ANY(${param_index})
                )
            """)
            params.append(filter_user_ids)

        query = f"""
            SELECT c.id, extract(epoch from c.created_at) AS created_at_epoch
            FROM {self._get_table_name("conversations")} c
            WHERE {" AND ".join(conditions)}
        """

        conv_row = await self.connection_manager.fetchrow_query(query, params)
        if not conv_row:
            raise R2RException(
                status_code=404,
                message=f"Conversation {conversation_id} not found.",
            )

        # Retrieve messages in chronological order
        msg_query = f"""
            SELECT id, content, metadata
            FROM {self._get_table_name("messages")}
            WHERE conversation_id = $1
            ORDER BY created_at ASC
        """
        results = await self.connection_manager.fetch_query(
            msg_query, [conversation_id]
        )

        response_messages = []
        for row in results:
            try:
                # Parse the message content
                content_json = json.loads(row["content"])
                # Create a Message object with the parsed content
                message = Message(**content_json)
                # Create a MessageResponse
                response_messages.append(
                    MessageResponse(
                        id=row["id"],
                        message=message,
                        metadata=json.loads(row["metadata"]),
                    )
                )
            except Exception as e:
                # If there's an error parsing the message (e.g., due to version mismatch),
                # log it and create a fallback message
                logger.warning(f"Error parsing message {row['id']}: {str(e)}")
                fallback_content = content_json.get(
                    "content", "Message could not be loaded"
                )
                fallback_role = content_json.get("role", "assistant")

                # Create a basic fallback message
                fallback_message = Message(
                    role=fallback_role,
                    content=f"[Message format incompatible: {fallback_content}]",
                )

                response_messages.append(
                    MessageResponse(
                        id=row["id"],
                        message=fallback_message,
                        metadata=json.loads(row["metadata"]),
                    )
                )

        return response_messages

    async def update_conversation(
        self, conversation_id: UUID, name: str
    ) -> ConversationResponse:
        try:
            # Check if conversation exists
            conv_query = f"SELECT 1 FROM {self._get_table_name('conversations')} WHERE id = $1"
            conv_row = await self.connection_manager.fetchrow_query(
                conv_query, [conversation_id]
            )
            if not conv_row:
                raise R2RException(
                    status_code=404,
                    message=f"Conversation {conversation_id} not found.",
                )

            update_query = f"""
            UPDATE {self._get_table_name("conversations")}
            SET name = $1 WHERE id = $2
            RETURNING user_id, extract(epoch from created_at) as created_at_epoch
            """
            updated_row = await self.connection_manager.fetchrow_query(
                update_query, [name, conversation_id]
            )
            return ConversationResponse(
                id=conversation_id,
                created_at=updated_row["created_at_epoch"],
                user_id=updated_row["user_id"] or None,
                name=name,
            )
        except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"Failed to update conversation: {str(e)}",
            ) from e

    async def delete_conversation(
        self,
        conversation_id: UUID,
        filter_user_ids: Optional[list[UUID]] = None,
    ) -> None:
        conditions = ["c.id = $1"]
        params: list = [conversation_id]

        if filter_user_ids:
            param_index = 2
            conditions.append(f"""
                c.user_id IN (
                    SELECT id
                    FROM {self.project_name}.users
                    WHERE id = ANY(${param_index})
                )
            """)
            params.append(filter_user_ids)

        conv_query = f"""
            SELECT 1
            FROM {self._get_table_name("conversations")} c
            WHERE {" AND ".join(conditions)}
        """
        conv_row = await self.connection_manager.fetchrow_query(
            conv_query, params
        )
        if not conv_row:
            raise R2RException(
                status_code=404,
                message=f"Conversation {conversation_id} not found.",
            )

        # Delete all messages
        del_messages_query = f"DELETE FROM {self._get_table_name('messages')} WHERE conversation_id = $1"
        await self.connection_manager.execute_query(
            del_messages_query, [conversation_id]
        )

        # Delete conversation
        del_conv_query = f"DELETE FROM {self._get_table_name('conversations')} WHERE id = $1"
        await self.connection_manager.execute_query(
            del_conv_query, [conversation_id]
        )

    async def export_conversations_to_csv(
        self,
        columns: Optional[list[str]] = None,
        filters: Optional[dict] = None,
        include_header: bool = True,
    ) -> tuple[str, IO]:
        """Creates a CSV file from the PostgreSQL data and returns the path to
        the temp file."""
        valid_columns = {
            "id",
            "user_id",
            "created_at",
            "name",
        }

        if not columns:
            columns = list(valid_columns)
        elif invalid_cols := set(columns) - valid_columns:
            raise ValueError(f"Invalid columns: {invalid_cols}")

        select_stmt = f"""
            SELECT
                id::text,
                user_id::text,
                to_char(created_at, 'YYYY-MM-DD HH24:MI:SS') AS created_at,
                name
            FROM {self._get_table_name("conversations")}
        """

        conditions = []
        params: list[Any] = []
        param_index = 1

        if filters:
            for field, value in filters.items():
                if field not in valid_columns:
                    continue

                if isinstance(value, dict):
                    for op, val in value.items():
                        if op == "$eq":
                            conditions.append(f"{field} = ${param_index}")
                            params.append(val)
                            param_index += 1
                        elif op == "$gt":
                            conditions.append(f"{field} > ${param_index}")
                            params.append(val)
                            param_index += 1
                        elif op == "$lt":
                            conditions.append(f"{field} < ${param_index}")
                            params.append(val)
                            param_index += 1
                else:
                    # Direct equality
                    conditions.append(f"{field} = ${param_index}")
                    params.append(value)
                    param_index += 1

        if conditions:
            select_stmt = f"{select_stmt} WHERE {' AND '.join(conditions)}"

        select_stmt = f"{select_stmt} ORDER BY created_at DESC"

        temp_file = None
        try:
            temp_file = tempfile.NamedTemporaryFile(
                mode="w", delete=True, suffix=".csv"
            )
            writer = csv.writer(temp_file, quoting=csv.QUOTE_ALL)

            async with self.connection_manager.pool.get_connection() as conn:  # type: ignore
                async with conn.transaction():
                    cursor = await conn.cursor(select_stmt, *params)

                    if include_header:
                        writer.writerow(columns)

                    chunk_size = 1000
                    while True:
                        rows = await cursor.fetch(chunk_size)
                        if not rows:
                            break
                        for row in rows:
                            row_dict = {
                                "id": row[0],
                                "user_id": row[1],
                                "created_at": row[2],
                                "name": row[3],
                            }
                            writer.writerow([row_dict[col] for col in columns])

            temp_file.flush()
            return temp_file.name, temp_file

        except Exception as e:
            if temp_file:
                temp_file.close()
            raise HTTPException(
                status_code=500,
                detail=f"Failed to export data: {str(e)}",
            ) from e

    async def export_messages_to_csv(
        self,
        columns: Optional[list[str]] = None,
        filters: Optional[dict] = None,
        include_header: bool = True,
        handle_images: str = "metadata_only",  # Options: "full", "metadata_only", "exclude"
    ) -> tuple[str, IO]:
        """
        Creates a CSV file from the PostgreSQL data and returns the path to the temp file.

        Args:
            columns: List of columns to include in export
            filters: Filter criteria for messages
            include_header: Whether to include header row
            handle_images: How to handle image data in exports:
                - "full": Include complete image data (warning: may create large files)
                - "metadata_only": Replace image data with metadata only
                - "exclude": Remove image data completely
        """
        valid_columns = {
            "id",
            "conversation_id",
            "parent_id",
            "content",
            "metadata",
            "created_at",
            "has_image",  # New virtual column to indicate image presence
        }

        if not columns:
            columns = list(valid_columns - {"has_image"})
        elif invalid_cols := set(columns) - valid_columns:
            raise ValueError(f"Invalid columns: {invalid_cols}")

        # Add virtual column for image presence
        virtual_columns = []
        has_image_column = False

        if "has_image" in columns:
            virtual_columns.append(
                "(content->>'image_url' IS NOT NULL OR content->>'image_data' IS NOT NULL) as has_image"
            )
            columns.remove("has_image")
            has_image_column = True

        select_stmt = f"""
            SELECT
                id::text,
                conversation_id::text,
                parent_id::text,
                content::text,
                metadata::text,
                to_char(created_at, 'YYYY-MM-DD HH24:MI:SS') AS created_at
                {", " + ", ".join(virtual_columns) if virtual_columns else ""}
            FROM {self._get_table_name("messages")}
        """

        # Keep existing filter conditions setup
        conditions = []
        params: list[Any] = []
        param_index = 1

        if filters:
            for field, value in filters.items():
                if field not in valid_columns or field == "has_image":
                    continue

                if isinstance(value, dict):
                    for op, val in value.items():
                        if op == "$eq":
                            conditions.append(f"{field} = ${param_index}")
                            params.append(val)
                            param_index += 1
                        elif op == "$gt":
                            conditions.append(f"{field} > ${param_index}")
                            params.append(val)
                            param_index += 1
                        elif op == "$lt":
                            conditions.append(f"{field} < ${param_index}")
                            params.append(val)
                            param_index += 1
                else:
                    conditions.append(f"{field} = ${param_index}")
                    params.append(value)
                    param_index += 1

        # Special filter for has_image
        if filters and "has_image" in filters:
            if filters["has_image"]:
                conditions.append(
                    "(content->>'image_url' IS NOT NULL OR content->>'image_data' IS NOT NULL)"
                )

        if conditions:
            select_stmt = f"{select_stmt} WHERE {' AND '.join(conditions)}"

        select_stmt = f"{select_stmt} ORDER BY created_at DESC"

        temp_file = None
        try:
            temp_file = tempfile.NamedTemporaryFile(
                mode="w", delete=True, suffix=".csv"
            )
            writer = csv.writer(temp_file, quoting=csv.QUOTE_ALL)

            # Prepare export columns
            export_columns = list(columns)
            if has_image_column:
                export_columns.append("has_image")

            if include_header:
                writer.writerow(export_columns)

            async with self.connection_manager.pool.get_connection() as conn:  # type: ignore
                async with conn.transaction():
                    cursor = await conn.cursor(select_stmt, *params)

                    chunk_size = 1000
                    while True:
                        rows = await cursor.fetch(chunk_size)
                        if not rows:
                            break
                        for row in rows:
                            row_dict = {
                                "id": row[0],
                                "conversation_id": row[1],
                                "parent_id": row[2],
                                "content": row[3],
                                "metadata": row[4],
                                "created_at": row[5],
                            }

                            # Add virtual column if present
                            if has_image_column:
                                row_dict["has_image"] = (
                                    "true" if row[6] else "false"
                                )

                            # Process image data based on handle_images setting
                            if (
                                "content" in columns
                                and handle_images != "full"
                            ):
                                try:
                                    content_json = json.loads(
                                        row_dict["content"]
                                    )

                                    if (
                                        "image_data" in content_json
                                        and content_json["image_data"]
                                    ):
                                        media_type = content_json[
                                            "image_data"
                                        ].get("media_type", "image/jpeg")

                                        if handle_images == "metadata_only":
                                            content_json["image_data"] = {
                                                "media_type": media_type,
                                                "data": "[BASE64_DATA_EXCLUDED_FROM_EXPORT]",
                                            }
                                        elif handle_images == "exclude":
                                            content_json.pop(
                                                "image_data", None
                                            )

                                    row_dict["content"] = json.dumps(
                                        content_json
                                    )
                                except (json.JSONDecodeError, TypeError) as e:
                                    logger.warning(
                                        f"Error processing message content for export: {e}"
                                    )

                            writer.writerow(
                                [row_dict[col] for col in export_columns]
                            )

            temp_file.flush()
            return temp_file.name, temp_file

        except Exception as e:
            if temp_file:
                temp_file.close()
            raise HTTPException(
                status_code=500,
                detail=f"Failed to export data: {str(e)}",
            ) from e