"""
CRUD ENDPOINTS FOR SEARCH TOOLS
"""
from datetime import datetime
from typing import Any, Dict, List, Union

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

from litellm._logging import verbose_proxy_logger
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


@router.get(
    "/search_tools/list",
    tags=["Search Tools"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=ListSearchToolsResponse,
)
async def list_search_tools():
    """
    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 search_tool in config_search_tools:
            tool_name = search_tool.get("search_tool_name")
            if tool_name:
                litellm_params_dict = dict(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=None,
                        search_tool_name=tool_name,
                        litellm_params=masked_litellm_params_dict,
                        search_tool_info=search_tool.get("search_tool_info"),
                        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 search_tool in search_tools_from_db:
            litellm_params_dict = dict(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=search_tool.get("search_tool_id"),
                    search_tool_name=search_tool.get("search_tool_name", ""),
                    litellm_params=masked_litellm_params_dict,
                    search_tool_info=search_tool.get("search_tool_info"),
                    created_at=_convert_datetime_to_str(search_tool.get("created_at")),
                    updated_at=_convert_datetime_to_str(search_tool.get("updated_at")),
                    is_from_config=False,
                )
            )

        return ListSearchToolsResponse(search_tools=search_tool_configs)
    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))