"""
Search Tool Registry for managing search tool configurations.
"""

from datetime import datetime, timezone
from typing import List, Optional

from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import SearchToolsRepository
from litellm.types.search import SearchTool


class SearchToolRegistry:
    """
    Handles adding, removing, and getting search tools in DB + in memory.
    """

    def __init__(self):
        pass

    @staticmethod
    def _convert_prisma_to_dict(prisma_obj) -> dict:
        """
        Convert Prisma result to dict with datetime objects as ISO format strings.

        Args:
            prisma_obj: Prisma model instance

        Returns:
            Dict with datetime fields converted to ISO strings
        """
        result = dict(prisma_obj)
        # Convert datetime objects to ISO format strings
        if "created_at" in result and result["created_at"]:
            result["created_at"] = result["created_at"].isoformat()
        if "updated_at" in result and result["updated_at"]:
            result["updated_at"] = result["updated_at"].isoformat()
        return result

    ###########################################################
    ########### DB management helpers for search tools ########
    ###########################################################

    async def add_search_tool_to_db(
        self, search_tool: SearchTool, prisma_client: PrismaClient
    ):
        """
        Add a search tool to the database.

        Args:
            search_tool: Search tool configuration
            prisma_client: Prisma client instance

        Returns:
            Dict with created search tool data
        """
        try:
            search_tool_name = search_tool.get("search_tool_name")
            litellm_params: str = safe_dumps(
                dict(search_tool.get("litellm_params", {}))
            )
            search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {}))

            # Create search tool in DB
            created_search_tool = await SearchToolsRepository(
                prisma_client
            ).table.create(
                data={
                    "search_tool_name": search_tool_name,
                    "litellm_params": litellm_params,
                    "search_tool_info": search_tool_info,
                    "created_at": datetime.now(timezone.utc),
                    "updated_at": datetime.now(timezone.utc),
                }
            )

            # Add search_tool_id to the returned search tool object
            search_tool_dict = dict(search_tool)
            search_tool_dict["search_tool_id"] = created_search_tool.search_tool_id
            search_tool_dict["created_at"] = created_search_tool.created_at.isoformat()
            search_tool_dict["updated_at"] = created_search_tool.updated_at.isoformat()

            return search_tool_dict
        except Exception as e:
            verbose_proxy_logger.exception(f"Error adding search tool to DB: {str(e)}")
            raise Exception(f"Error adding search tool to DB: {str(e)}")

    async def delete_search_tool_from_db(
        self, search_tool_id: str, prisma_client: PrismaClient
    ):
        """
        Delete a search tool from the database.

        Args:
            search_tool_id: ID of search tool to delete
            prisma_client: Prisma client instance

        Returns:
            Dict with success message
        """
        try:
            # Get search tool before deletion for response
            existing_tool = await SearchToolsRepository(
                prisma_client
            ).table.find_unique(where={"search_tool_id": search_tool_id})

            if not existing_tool:
                raise Exception(f"Search tool with ID {search_tool_id} not found")

            # Delete from DB
            await SearchToolsRepository(prisma_client).table.delete(
                where={"search_tool_id": search_tool_id}
            )

            return {
                "message": f"Search tool {search_tool_id} deleted successfully",
                "search_tool_name": existing_tool.search_tool_name,
            }
        except Exception as e:
            verbose_proxy_logger.exception(
                f"Error deleting search tool from DB: {str(e)}"
            )
            raise Exception(f"Error deleting search tool from DB: {str(e)}")

    async def update_search_tool_in_db(
        self, search_tool_id: str, search_tool: SearchTool, prisma_client: PrismaClient
    ):
        """
        Update a search tool in the database.

        Args:
            search_tool_id: ID of search tool to update
            search_tool: Updated search tool configuration
            prisma_client: Prisma client instance

        Returns:
            Dict with updated search tool data
        """
        try:
            search_tool_name = search_tool.get("search_tool_name")
            litellm_params: str = safe_dumps(
                dict(search_tool.get("litellm_params", {}))
            )
            search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {}))

            # Update in DB
            updated_search_tool = await SearchToolsRepository(
                prisma_client
            ).table.update(
                where={"search_tool_id": search_tool_id},
                data={
                    "search_tool_name": search_tool_name,
                    "litellm_params": litellm_params,
                    "search_tool_info": search_tool_info,
                    "updated_at": datetime.now(timezone.utc),
                },
            )

            # Convert to dict with ISO formatted datetimes
            return self._convert_prisma_to_dict(updated_search_tool)
        except Exception as e:
            verbose_proxy_logger.exception(
                f"Error updating search tool in DB: {str(e)}"
            )
            raise Exception(f"Error updating search tool in DB: {str(e)}")

    @staticmethod
    async def get_all_search_tools_from_db(
        prisma_client: PrismaClient,
    ) -> List[SearchTool]:
        """
        Get all search tools from the database.

        Args:
            prisma_client: Prisma client instance

        Returns:
            List of search tool configurations
        """
        try:
            search_tools_from_db = await call_with_db_reconnect_retry(
                prisma_client,
                lambda: SearchToolsRepository(prisma_client).table.find_many(
                    order={"created_at": "desc"},
                ),
                reason="get_all_search_tools_from_db_lookup_failure",
            )

            search_tools: List[SearchTool] = []
            for search_tool in search_tools_from_db:
                # Convert Prisma result to dict with ISO formatted datetimes
                search_tool_dict = SearchToolRegistry._convert_prisma_to_dict(
                    search_tool
                )
                search_tools.append(SearchTool(**search_tool_dict))  # type: ignore

            return search_tools
        except Exception as e:
            verbose_proxy_logger.exception(
                f"Error getting search tools from DB: {str(e)}"
            )
            raise Exception(f"Error getting search tools from DB: {str(e)}")

    async def get_search_tool_by_id_from_db(
        self, search_tool_id: str, prisma_client: PrismaClient
    ) -> Optional[SearchTool]:
        """
        Get a search tool by its ID from the database.

        Args:
            search_tool_id: ID of search tool to retrieve
            prisma_client: Prisma client instance

        Returns:
            Search tool configuration or None if not found
        """
        try:
            search_tool = await SearchToolsRepository(prisma_client).table.find_unique(
                where={"search_tool_id": search_tool_id}
            )

            if not search_tool:
                return None

            # Convert Prisma result to dict with ISO formatted datetimes
            search_tool_dict = self._convert_prisma_to_dict(search_tool)
            return SearchTool(**search_tool_dict)  # type: ignore
        except Exception as e:
            verbose_proxy_logger.exception(
                f"Error getting search tool from DB: {str(e)}"
            )
            raise Exception(f"Error getting search tool from DB: {str(e)}")

    async def get_search_tool_by_name_from_db(
        self, search_tool_name: str, prisma_client: PrismaClient
    ) -> Optional[SearchTool]:
        """
        Get a search tool by its name from the database.

        Args:
            search_tool_name: Name of search tool to retrieve
            prisma_client: Prisma client instance

        Returns:
            Search tool configuration or None if not found
        """
        try:
            search_tool = await SearchToolsRepository(prisma_client).table.find_unique(
                where={"search_tool_name": search_tool_name}
            )

            if not search_tool:
                return None

            # Convert Prisma result to dict with ISO formatted datetimes
            search_tool_dict = self._convert_prisma_to_dict(search_tool)
            return SearchTool(**search_tool_dict)  # type: ignore
        except Exception as e:
            verbose_proxy_logger.exception(
                f"Error getting search tool from DB: {str(e)}"
            )
            raise Exception(f"Error getting search tool from DB: {str(e)}")
