from abc import ABC, abstractmethod
from enum import StrEnum
from typing import Generic, TypeVar

from pydantic import BaseModel, Field
from unique_toolkit.agentic.tools.config import get_configuration_dict

from unique_web_search.services.helpers import (
    clean_model_title_generator,
)
from unique_web_search.services.search_engine.schema import (
    WebSearchResult,
)


class SearchEngineType(StrEnum):
    GOOGLE = "Google"
    JINA = "Jina"
    FIRECRAWL = "Firecrawl"
    TAVILY = "Tavily"
    BRAVE = "Brave"
    BING = "Bing"
    DUCKDUCKGO = "DuckDuckGo"
    VERTEXAI = "VertexAI"
    CUSTOM_API = "CustomAPI"


T = TypeVar("T", bound=SearchEngineType)


class BaseSearchEngineConfig(BaseModel, Generic[T]):
    model_config = get_configuration_dict(
        model_title_generator=clean_model_title_generator
    )
    search_engine_name: T

    fetch_size: int = Field(
        default=5,
        description="Number of search results to fetch",
    )


SearchEngineConfig = TypeVar(
    "SearchEngineConfig",
    bound=BaseSearchEngineConfig,
)


class SearchEngine(ABC, Generic[SearchEngineConfig]):
    """Base class for the search engine. It provides the common methods to search the web. This allows to use different search engines easily.

    Abstract Methods:
        _extract_urls: Clean the response from the search engine.
        _get_request_params: Get the request parameters.
        requires_scraping: Whether the search engine requires scraping.

    Args:
        assistant: The assistant configuration.
        chat_manager: The chat manager.

    """

    def __init__(self, config: SearchEngineConfig):
        self.config = config

    @property
    @abstractmethod
    def requires_scraping(self) -> bool:
        """Whether the search engine requires scraping."""

    @abstractmethod
    async def search(self, query: str, **kwargs) -> list[WebSearchResult]:
        """Search the web for the given query using the search engine."""
