"""
CRUD ENDPOINTS FOR SEARCH TOOLS
"""

from datetime import datetime
from typing import Any, Dict, List, Optional, Union

from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel

from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
    LiteLLM_TeamTable,
    LitellmUserRoles,
    UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.search_endpoints.search_tool_registry import SearchToolRegistry
from litellm.types.search import (
    ListSearchToolsResponse,
    SearchTool,
    SearchToolInfoResponse,
)
from litellm.types.utils import SearchProviders

#### SEARCH TOOLS ENDPOINTS ####

router = APIRouter()
SEARCH_TOOL_REGISTRY = SearchToolRegistry()


def _convert_datetime_to_str(value: Union[datetime, str, None]) -> Union[str, None]:
    """
    Convert datetime object to ISO format string.

    Args:
        value: datetime object, string, or None

    Returns:
        ISO format string or original value if already string or None
    """
    if value is None:
        return None
    if isinstance(value, datetime):
        return value.isoformat()
    return value


async def _filter_visible_search_tools(
    search_tools: List[SearchToolInfoResponse],
    user_api_key_dict: UserAPIKeyAuth,
) -> List[SearchToolInfoResponse]:
    """
    Drop search tools the caller is not authorized to invoke, applying the same
    key/team object_permission allowlists enforced on /search. Admins see all tools.
    """
    if user_api_key_dict.user_role in (
        LitellmUserRoles.PROXY_ADMIN,
        LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
    ):
        return search_tools

    from litellm.proxy.auth.auth_checks import (
        can_user_view_search_tool,
        get_team_object,
    )
    from litellm.proxy.proxy_server import (
        prisma_client,
        proxy_logging_obj,
        user_api_key_cache,
    )

    team_object: Optional[LiteLLM_TeamTable] = None
    if user_api_key_dict.team_id:
        team_object = await get_team_object(
            team_id=user_api_key_dict.team_id,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=user_api_key_dict.parent_otel_span,
            proxy_logging_obj=proxy_logging_obj,
        )

    visible: List[SearchToolInfoResponse] = []
    for tool in search_tools:
        tool_name = tool.get("search_tool_name")
        if tool_name and await can_user_view_search_tool(
            search_tool_name=tool_name,
            valid_token=user_api_key_dict,
            team_object=team_object,
        ):
            visible.append(tool)
    return visible


@router.get(
    "/search_tools/list",
    tags=["Search Tools"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=ListSearchToolsResponse,
)
async def list_search_tools(
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    List all search tools that are available in the database and config file.

    Example Request:
    ```bash
    curl -X GET "http://localhost:4000/search_tools/list" -H "Authorization: Bearer <your_api_key>"
    ```

    Example Response:
    ```json
    {
        "search_tools": [
            {
                "search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
                "search_tool_name": "litellm-search",
                "litellm_params": {
                    "search_provider": "perplexity",
                    "api_key": "sk-***",
                    "api_base": "https://api.perplexity.ai"
                },
                "search_tool_info": {
                    "description": "Perplexity search tool"
                },
                "created_at": "2023-11-09T12:34:56.789Z",
                "updated_at": "2023-11-09T12:34:56.789Z",
                "is_from_config": false
            },
            {
                "search_tool_name": "config-search-tool",
                "litellm_params": {
                    "search_provider": "tavily",
                    "api_key": "tvly-***"
                },
                "is_from_config": true
            }
        ]
    }
    ```
    """
    from litellm.litellm_core_utils.litellm_logging import _get_masked_values
    from litellm.proxy.proxy_server import prisma_client, proxy_config

    if prisma_client is None:
        raise HTTPException(status_code=500, detail="Prisma client not initialized")

    try:
        search_tools_from_db = await SEARCH_TOOL_REGISTRY.get_all_search_tools_from_db(
            prisma_client=prisma_client
        )

        db_tool_names = {tool.get("search_tool_name") for tool in search_tools_from_db}

        search_tool_configs: List[SearchToolInfoResponse] = []

        config_search_tools = []

        try:
            config = await proxy_config.get_config()
            parsed_tools = proxy_config.parse_search_tools(config)
            if parsed_tools:
                config_search_tools = parsed_tools
        except Exception as e:
            verbose_proxy_logger.debug(
                f"Could not get config-defined search tools: {e}"
            )

        for config_search_tool in config_search_tools:
            tool_name = config_search_tool.get("search_tool_name")
            if tool_name:
                litellm_params_dict = dict(config_search_tool.get("litellm_params", {}))
                masked_litellm_params_dict = _get_masked_values(
                    litellm_params_dict,
                    unmasked_length=4,
                    number_of_asterisks=4,
                )
                config_tool_info = config_search_tool.get("search_tool_info")

                search_tool_configs.append(
                    SearchToolInfoResponse(
                        search_tool_id=None,
                        search_tool_name=tool_name,
                        litellm_params=masked_litellm_params_dict,
                        search_tool_info=(
                            dict(config_tool_info) if config_tool_info else None
                        ),
                        created_at=None,
                        updated_at=None,
                        is_from_config=True,
                    )
                )

        search_tool_configs = [
            tool
            for tool in search_tool_configs
            if tool.get("search_tool_name") not in db_tool_names
        ]

        for db_search_tool in search_tools_from_db:
            litellm_params_dict = dict(db_search_tool.get("litellm_params", {}))
            masked_litellm_params_dict = _get_masked_values(
                litellm_params_dict,
                unmasked_length=4,
                number_of_asterisks=4,
            )

            search_tool_configs.append(
                SearchToolInfoResponse(
                    search_tool_id=db_search_tool.get("search_tool_id"),
                    search_tool_name=db_search_tool.get("search_tool_name", ""),
                    litellm_params=masked_litellm_params_dict,
                    search_tool_info=db_search_tool.get("search_tool_info"),
                    created_at=_convert_datetime_to_str(
                        db_search_tool.get("created_at")
                    ),
                    updated_at=_convert_datetime_to_str(
                        db_search_tool.get("updated_at")
                    ),
                    is_from_config=False,
                )
            )

        visible_search_tools = await _filter_visible_search_tools(
            search_tool_configs, user_api_key_dict
        )

        return ListSearchToolsResponse(search_tools=visible_search_tools)
    except Exception as e:
        verbose_proxy_logger.exception(f"Error getting search tools: {e}")
        raise HTTPException(status_code=500, detail=str(e))


class CreateSearchToolRequest(BaseModel):
    search_tool: SearchTool


@router.post(
    "/search_tools",
    tags=["Search Tools"],
    dependencies=[Depends(user_api_key_auth)],
)
async def create_search_tool(request: CreateSearchToolRequest):
    """
    Create a new search tool.

    Example Request:
    ```bash
    curl -X POST "http://localhost:4000/search_tools" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json" \\
        -d '{
            "search_tool": {
                "search_tool_name": "litellm-search",
                "litellm_params": {
                    "search_provider": "perplexity",
                    "api_key": "sk-..."
                },
                "search_tool_info": {
                    "description": "Perplexity search tool"
                }
            }
        }'
    ```

    Example Response:
    ```json
    {
        "search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
        "search_tool_name": "litellm-search",
        "litellm_params": {
            "search_provider": "perplexity",
            "api_key": "sk-..."
        },
        "search_tool_info": {
            "description": "Perplexity search tool"
        },
        "created_at": "2023-11-09T12:34:56.789Z",
        "updated_at": "2023-11-09T12:34:56.789Z"
    }
    ```
    """
    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(status_code=500, detail="Prisma client not initialized")

    try:
        result = await SEARCH_TOOL_REGISTRY.add_search_tool_to_db(
            search_tool=request.search_tool, prisma_client=prisma_client
        )

        verbose_proxy_logger.debug(
            f"Successfully added search tool '{result.get('search_tool_name')}' to database. "
            f"Router will be updated by the cron job."
        )

        return result
    except Exception as e:
        verbose_proxy_logger.exception(f"Error adding search tool to db: {e}")
        raise HTTPException(status_code=500, detail=str(e))


class UpdateSearchToolRequest(BaseModel):
    search_tool: SearchTool


@router.put(
    "/search_tools/{search_tool_id}",
    tags=["Search Tools"],
    dependencies=[Depends(user_api_key_auth)],
)
async def update_search_tool(search_tool_id: str, request: UpdateSearchToolRequest):
    """
    Update an existing search tool.

    Example Request:
    ```bash
    curl -X PUT "http://localhost:4000/search_tools/123e4567-e89b-12d3-a456-426614174000" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json" \\
        -d '{
            "search_tool": {
                "search_tool_name": "updated-search",
                "litellm_params": {
                    "search_provider": "perplexity",
                    "api_key": "sk-new-key"
                },
                "search_tool_info": {
                    "description": "Updated search tool"
                }
            }
        }'
    ```

    Example Response:
    ```json
    {
        "search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
        "search_tool_name": "updated-search",
        "litellm_params": {
            "search_provider": "perplexity",
            "api_key": "sk-new-key"
        },
        "search_tool_info": {
            "description": "Updated search tool"
        },
        "created_at": "2023-11-09T12:34:56.789Z",
        "updated_at": "2023-11-09T13:45:12.345Z"
    }
    ```
    """
    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(status_code=500, detail="Prisma client not initialized")

    try:
        # Check if search tool exists
        existing_tool = await SEARCH_TOOL_REGISTRY.get_search_tool_by_id_from_db(
            search_tool_id=search_tool_id, prisma_client=prisma_client
        )

        if existing_tool is None:
            raise HTTPException(
                status_code=404,
                detail=f"Search tool with ID {search_tool_id} not found",
            )

        result = await SEARCH_TOOL_REGISTRY.update_search_tool_in_db(
            search_tool_id=search_tool_id,
            search_tool=request.search_tool,
            prisma_client=prisma_client,
        )

        verbose_proxy_logger.debug(
            f"Successfully updated search tool '{result.get('search_tool_name')}' in database. "
            f"Router will be updated by the cron job."
        )

        return result
    except HTTPException as e:
        raise e
    except Exception as e:
        verbose_proxy_logger.exception(f"Error updating search tool: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.delete(
    "/search_tools/{search_tool_id}",
    tags=["Search Tools"],
    dependencies=[Depends(user_api_key_auth)],
)
async def delete_search_tool(search_tool_id: str):
    """
    Delete a search tool.

    Example Request:
    ```bash
    curl -X DELETE "http://localhost:4000/search_tools/123e4567-e89b-12d3-a456-426614174000" \\
        -H "Authorization: Bearer <your_api_key>"
    ```

    Example Response:
    ```json
    {
        "message": "Search tool 123e4567-e89b-12d3-a456-426614174000 deleted successfully",
        "search_tool_name": "litellm-search"
    }
    ```
    """
    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(status_code=500, detail="Prisma client not initialized")

    try:
        # Check if search tool exists
        existing_tool = await SEARCH_TOOL_REGISTRY.get_search_tool_by_id_from_db(
            search_tool_id=search_tool_id, prisma_client=prisma_client
        )

        if existing_tool is None:
            raise HTTPException(
                status_code=404,
                detail=f"Search tool with ID {search_tool_id} not found",
            )

        result = await SEARCH_TOOL_REGISTRY.delete_search_tool_from_db(
            search_tool_id=search_tool_id, prisma_client=prisma_client
        )

        verbose_proxy_logger.debug(
            "Successfully deleted search tool from database. "
            "Router will be updated by the cron job."
        )

        return result
    except HTTPException as e:
        raise e
    except Exception as e:
        verbose_proxy_logger.exception(f"Error deleting search tool: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get(
    "/search_tools/{search_tool_id}",
    tags=["Search Tools"],
    dependencies=[Depends(user_api_key_auth)],
)
async def get_search_tool_info(search_tool_id: str):
    """
    Get detailed information about a specific search tool by ID.

    Example Request:
    ```bash
    curl -X GET "http://localhost:4000/search_tools/123e4567-e89b-12d3-a456-426614174000" \\
        -H "Authorization: Bearer <your_api_key>"
    ```

    Example Response:
    ```json
    {
        "search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
        "search_tool_name": "litellm-search",
        "litellm_params": {
            "search_provider": "perplexity",
            "api_key": "sk-***"
        },
        "search_tool_info": {
            "description": "Perplexity search tool"
        },
        "created_at": "2023-11-09T12:34:56.789Z",
        "updated_at": "2023-11-09T12:34:56.789Z"
    }
    ```
    """
    from litellm.litellm_core_utils.litellm_logging import _get_masked_values
    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(status_code=500, detail="Prisma client not initialized")

    try:
        result = await SEARCH_TOOL_REGISTRY.get_search_tool_by_id_from_db(
            search_tool_id=search_tool_id, prisma_client=prisma_client
        )

        if result is None:
            raise HTTPException(
                status_code=404,
                detail=f"Search tool with ID {search_tool_id} not found",
            )

        # Mask sensitive data
        litellm_params_dict = dict(result.get("litellm_params", {}))
        masked_litellm_params_dict = _get_masked_values(
            litellm_params_dict,
            unmasked_length=4,
            number_of_asterisks=4,
        )

        return SearchToolInfoResponse(
            search_tool_id=result.get("search_tool_id"),
            search_tool_name=result.get("search_tool_name", ""),
            litellm_params=masked_litellm_params_dict,
            search_tool_info=result.get("search_tool_info"),
            created_at=_convert_datetime_to_str(result.get("created_at")),
            updated_at=_convert_datetime_to_str(result.get("updated_at")),
            is_from_config=False,  # This endpoint only returns DB tools
        )
    except HTTPException as e:
        raise e
    except Exception as e:
        verbose_proxy_logger.exception(f"Error getting search tool info: {e}")
        raise HTTPException(status_code=500, detail=str(e))


class TestSearchToolConnectionRequest(BaseModel):
    litellm_params: Dict[str, Any]


@router.post(
    "/search_tools/test_connection",
    tags=["Search Tools"],
    dependencies=[Depends(user_api_key_auth)],
)
async def test_search_tool_connection(request: TestSearchToolConnectionRequest):
    """
    Test connection to a search provider with the given configuration.
    
    Makes a simple test search query to verify the API key and configuration are valid.

    Example Request:
    ```bash
    curl -X POST "http://localhost:4000/search_tools/test_connection" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json" \\
        -d '{
            "litellm_params": {
                "search_provider": "perplexity",
                "api_key": "sk-..."
            }
        }'
    ```

    Example Response (Success):
    ```json
    {
        "status": "success",
        "message": "Successfully connected to perplexity search provider",
        "test_query": "test",
        "results_count": 5
    }
    ```

    Example Response (Failure):
    ```json
    {
        "status": "error",
        "message": "Authentication failed: Invalid API key",
        "error_type": "AuthenticationError"
    }
    ```
    """
    try:
        from litellm.search import asearch

        # Extract params from request
        litellm_params = request.litellm_params
        search_provider = litellm_params.get("search_provider")
        api_key = litellm_params.get("api_key")
        api_base = litellm_params.get("api_base")

        if not search_provider:
            raise HTTPException(
                status_code=400, detail="search_provider is required in litellm_params"
            )

        verbose_proxy_logger.debug(
            f"Testing connection to search provider: {search_provider}"
        )

        # Make a simple test search query with max_results=1 to minimize cost
        test_query = "test"
        response = await asearch(
            query=test_query,
            search_provider=search_provider,
            api_key=api_key,
            api_base=api_base,
            max_results=1,  # Minimize results to reduce cost
            timeout=10.0,  # 10 second timeout for test
        )

        verbose_proxy_logger.debug(
            f"Successfully tested connection to {search_provider} search provider"
        )

        return {
            "status": "success",
            "message": f"Successfully connected to {search_provider} search provider",
            "test_query": test_query,
            "results_count": (
                len(response.results) if response and response.results else 0
            ),
        }

    except Exception as e:
        error_message = str(e)
        error_type = type(e).__name__

        verbose_proxy_logger.exception(
            f"Failed to connect to search provider: {error_message}"
        )

        # Return error details in a structured format
        return {
            "status": "error",
            "message": error_message,
            "error_type": error_type,
        }


@router.get(
    "/search_tools/ui/available_providers",
    tags=["Search Tools"],
    dependencies=[Depends(user_api_key_auth)],
)
async def get_available_search_providers():
    """
    Get the list of available search providers with their configuration fields.
    
    Auto-discovers search providers and their UI-friendly names from transformation configs.

    Example Request:
    ```bash
    curl -X GET "http://localhost:4000/search_tools/ui/available_providers" \\
        -H "Authorization: Bearer <your_api_key>"
    ```

    Example Response:
    ```json
    {
        "providers": [
            {
                "provider_name": "perplexity",
                "ui_friendly_name": "Perplexity"
            },
            {
                "provider_name": "tavily",
                "ui_friendly_name": "Tavily"
            }
        ]
    }
    ```
    """
    try:
        from litellm.utils import ProviderConfigManager

        available_providers = []

        # Auto-discover providers from SearchProviders enum
        for provider in SearchProviders:
            try:
                # Get the config class for this provider
                config = ProviderConfigManager.get_provider_search_config(
                    provider=provider
                )

                if config is not None:
                    # Get the UI-friendly name from the config class
                    ui_name = config.ui_friendly_name()

                    available_providers.append(
                        {
                            "provider_name": provider.value,
                            "ui_friendly_name": ui_name,
                        }
                    )
            except Exception as e:
                verbose_proxy_logger.debug(
                    f"Could not get config for search provider {provider.value}: {e}"
                )
                continue

        return {"providers": available_providers}
    except Exception as e:
        verbose_proxy_logger.exception(f"Error getting available search providers: {e}")
        raise HTTPException(status_code=500, detail=str(e))
