import csv
import json
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


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,
    ) -> MessageResponse:
        # 1) Validate that conversation and parent exist, as before
        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) Convert the content & metadata to JSON strings in a safe manner
        message_id = uuid4()
        content_str = safe_dumps(content.model_dump())  # <-- use safe_dumps
        metadata_str = safe_dumps(metadata or {})

        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,
            )
            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]:
        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]
        )

        return [
            MessageResponse(
                id=row["id"],
                message=Message(**json.loads(row["content"])),
                metadata=json.loads(row["metadata"]),
            )
            for row in results
        ]

    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,
    ) -> tuple[str, IO]:
        """Creates a CSV file from the PostgreSQL data and returns the path to
        the temp file."""
        valid_columns = {
            "id",
            "conversation_id",
            "parent_id",
            "content",
            "metadata",
            "created_at",
        }

        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,
                conversation_id::text,
                parent_id::text,
                content::text,
                metadata::text,
                to_char(created_at, 'YYYY-MM-DD HH24:MI:SS') AS created_at
            FROM {self._get_table_name("messages")}
        """

        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],
                                "conversation_id": row[1],
                                "parent_id": row[2],
                                "content": row[3],
                                "metadata": row[4],
                                "created_at": row[5],
                            }
                            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
