# -*- coding: utf-8 -*-
# Copyright (c) Louis Brulé Naudet. All Rights Reserved.
# This software may be used and distributed according to the terms of License Agreement.
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import (
    IO,
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Type,
    Tuple,
    Union,
    Mapping,
    TypeVar,
    Callable,
    Optional,
    Sequence,
)

from ragoon._retrieval import Retriever
from ragoon._scrape import WebScraper
from ragoon._search import GoogleSearch


class WebRAG:
    def __init__(
        self,
        google_api_key: str,
        google_cx: str,
        completion_client,
        user_agent: Optional[str] = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/98.0.4758.102 Safari/537.36"
    ) -> None:
        """
        WebRAG class.

        This class facilitates retrieval-based querying and completion using various APIs.

        Parameters
        ----------
        google_api_key : str
            The API key for Google services.
        
        google_cx : str
            The custom search engine ID for Google Custom Search.
        
        completion_client: str
            The API client for the completion service (e.g., OpenAI's GPT-3).
        
        user_agent : str, optional
            The user agent string to be used in web requests. Default is a Chrome user agent.

        Attributes
        ----------
        web_scraper : WebScraper
            An instance of the WebScraper class for web scraping.
        
        retriever : Retriever
            An instance of the Retriever class for data retrieval.
        
        google_search : GoogleSearch
            An instance of the GoogleSearch class for Google searches.

        Examples
        --------
        # Initialize RAGoon instance
        >>> ragoon = RAGoon(
        >>>    google_api_key="your_google_api_key",
        >>>    google_cx="your_google_cx",
        >>>    completion_client=Groq(api_key="your_groq_api_key")
        >>> )

        >>> # Search and get results
        >>> query = "I want to do a left join in python polars"
        >>> results = ragoon.search(
        >>>    query=query,
        >>>    completion_model="Llama3-70b-8192",
        >>>    max_tokens=512,
        >>>    temperature=1,
        >>> )
        """
        self.web_search = GoogleSearch(
            developer_key=google_api_key, 
            cx=google_cx
        )
        self.retriever = Retriever(
            client=completion_client
        )
        self.web_scraper = WebScraper(
            user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/98.0.4758.102 Safari/537.36"
        )


    def search(
        self,
        query: str,
        completion_model: str,
        system_prompt: Optional[str] = """
        Given the user's input query, generate a concise and relevant Google search
        query that directly addresses the main intent of the user's question. The search query must
        be specifically tailored to retrieve results that can significantly enhance the context for a
        subsequent dialogue with an LLM. This approach will facilitate few-shot learning by providing
        rich, specific, and contextually relevant information. Please ensure that the response is
        well-formed and format it as a JSON object with a key named 'search_query'. This
        structured approach will help in assimilating the fetched results into an enhanced conversational
        model, contributing to a more nuanced and informed interaction.
        """,
        *args,
        **kargs
    ):
        """
        Search for information and perform completion.

        This method searches for information related to the given query
        and performs completion using the specified model. Additional
        parameters can be passed to the completion method.

        Parameters
        ----------
        query : str
            The search query.
        
        completion_model : str
            The name or identifier of the completion model to be used.
        
        *args
            Additional positional arguments to be passed to the completion method.
        
        **kwargs
            Additional keyword arguments to be passed to the completion method.

        Returns
        -------
        completion_data : dict
            A dictionary containing the generated completion data.

        """
        search_query = self.retriever.completion(
            model=completion_model,
            messages=[
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": query
                }
            ],
            *args
        )["search_query"]

        search_results = self.web_search.search(
            query=search_query
        )

        list_of_urls = []

        for result in search_results:
            list_of_urls.append(result.get("link"))

        results = self.web_scraper.parallel_scrape(
            urls=list_of_urls,
            element_selectors= [
                "main"
            ]
        )

        return results