"""
Brave Search AI Chat API Provider

This module provides a complete implementation of the Brave Search AI Chat API.
It supports both streaming and non-streaming responses using the real API endpoints.

Source: https://search.brave.com/ask
API Endpoint: https://search.brave.com/api/tap/v1

Features:
    - AI-generated responses based on web search
    - Streaming and non-streaming modes
    - Deep research capability
    - Proper error handling and session management

Generated by ReVibe.
Co-Authored-By: ReVibe <revibe@oevortex.dev>

STATUS: UNFINISHED
Issue: API endpoint returns 404 error on /new endpoint initialization
The reverse-engineered endpoints appear correct, but the actual API
may have changed or requires additional authentication/parameters.
Requires further investigation and testing against the live API.
"""

import asyncio
import base64
import json
import os
from typing import Any, AsyncIterator, Dict, Generator, List, Optional, Union, cast
from urllib.parse import quote

from curl_cffi.requests import AsyncSession

from webscout import exceptions
from webscout.AIbase import AISearch, SearchResponse


class BraveAI(AISearch):
    """
    A class to interact with Brave Search AI Chat API.

    BraveAI provides AI-generated responses based on web search results using the
    official Brave Search API. It supports both streaming and non-streaming responses.

    The API uses a symmetric_key for request authentication and requires various
    context parameters for proper operation.

    Basic Usage:
        >>> from webscout.Provider.AISEARCH import BraveAI
        >>> ai = BraveAI()
        >>> response = ai.search("What is Python?")
        >>> print(response)

    Streaming Usage:
        >>> ai = BraveAI()
        >>> for chunk in ai.search("Tell me about AI", stream=True):
        ...     print(chunk, end="", flush=True)

    Args:
        timeout (int, optional): Request timeout in seconds. Defaults to 30.
        proxies (dict, optional): Proxy configuration for requests. Defaults to None.
        language (str, optional): Language for search results. Defaults to "en".
        country (str, optional): Country code for localization. Defaults to "US".
        ui_lang (str, optional): UI language code. Defaults to "en-us".
        geoloc (str, optional): Geolocation in format "latitude,longitude".
                               Defaults to None (will use default).
        safesearch (str, optional): Safe search level. Defaults to "moderate".
                                   Options: "off", "moderate", "strict".
        units (str, optional): Units of measurement. Defaults to "metric".
        use_location (bool, optional): Whether to use geolocation. Defaults to True.
    """

    BASE_URL = "https://search.brave.com/api/tap/v1"

    def __init__(
        self,
        timeout: int = 30,
        proxies: Optional[dict] = None,
        language: str = "en",
        country: str = "US",
        ui_lang: str = "en-us",
        geoloc: Optional[str] = None,
        safesearch: str = "moderate",
        units: str = "metric",
        use_location: bool = True,
    ):
        """Initialize BraveAI with configuration parameters."""
        self.timeout = timeout
        self.proxies = proxies or {}
        self.language = language
        self.country = country.lower()
        self.ui_lang = ui_lang
        self.geoloc = geoloc
        self.safesearch = safesearch
        self.units = units
        self.use_location = "1" if use_location else "0"
        self.last_response = {}
        self.chat_id = None
        self.symmetric_key = None
        self.headers = {
            "accept": "application/json",
            "accept-language": "en-US,en;q=0.9",
            "user-agent": (
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                "(KHTML, like Gecko) Chrome/144.0.0.0 Safari/537.36"
            ),
            "sec-ch-ua": (
                '"Not(A:Brand";v="8", "Chromium";v="144", "Google Chrome";v="144"'
            ),
            "sec-ch-ua-mobile": "?0",
            "sec-ch-ua-platform": '"Windows"',
            "sec-fetch-dest": "empty",
            "sec-fetch-mode": "cors",
            "sec-fetch-site": "same-origin",
            "referer": "https://search.brave.com/ask",
        }

    def _generate_symmetric_key(self) -> str:
        """
        Generate a symmetric encryption key for API authentication.

        Creates a 256-bit AES-GCM key in JWK (JSON Web Key) format, base64 encoded.

        Returns:
            str: Base64-encoded JWK representation of the encryption key.
        """
        key = os.urandom(32)
        k = base64.urlsafe_b64encode(key).decode().rstrip("=")
        jwk = {
            "alg": "A256GCM",
            "ext": True,
            "k": k,
            "key_ops": ["encrypt", "decrypt"],
            "kty": "oct",
        }
        return base64.b64encode(
            json.dumps(jwk, separators=(",", ":")).encode()
        ).decode()

    def _build_context_params(
        self, enable_research: bool = False
    ) -> Dict[str, Union[str, int]]:
        """
        Build common context parameters for API requests.

        These parameters are required by the Brave Search API for proper
        context and localization of responses.

        Args:
            enable_research (bool): Whether to enable deep research mode.

        Returns:
            dict: Dictionary of context parameters.
        """
        params = {
            "language": self.language,
            "country": self.country,
            "ui_lang": self.ui_lang,
            "safesearch": self.safesearch,
            "force_safesearch": "0",
            "units_of_measurement": self.units,
            "use_location": self.use_location,
        }

        if self.geoloc:
            params["geoloc"] = self.geoloc

        return params

    def search(
        self,
        prompt: str,
        stream: bool = False,
        raw: bool = False,
        enable_research: bool = False,
        **kwargs: Any,
    ) -> Union[SearchResponse, Generator[Union[Dict[str, str], SearchResponse], None, None]]:
        """
        Search using the Brave Search AI and get AI-generated responses.

        This method sends a search query to Brave AI and returns the AI-generated
        response. It supports both streaming and non-streaming modes, as well as
        raw response format.

        Args:
            prompt (str): The search query or prompt to send to the API.
            stream (bool, optional): If True, yields response chunks as they arrive.
                                   If False, returns complete response. Defaults to False.
            raw (bool, optional): If True, returns raw response dictionaries with
                                'text' key. If False, returns Response objects.
                                Defaults to False.
            enable_research (bool, optional): If True, enables deep research mode
                                             for more comprehensive responses.
                                             Defaults to False.
            **kwargs: Additional keyword arguments (for compatibility).

        Returns:
            Union[SearchResponse, Generator[Union[Dict[str, str], SearchResponse], None, None]]:
                - If stream=False: Returns complete response as SearchResponse object
                - If stream=True: Yields response chunks as either Dict or SearchResponse

        Raises:
            APIConnectionError: If the API request fails.
            FailedToGenerateResponseError: If response parsing fails.

        Examples:
            Basic search:
            >>> ai = BraveAI()
            >>> response = ai.search("What is Python?")
            >>> print(response)

            Research-enabled search:
            >>> response = ai.search("What is Python?", enable_research=True)
            >>> print(response)

            Streaming response:
            >>> for chunk in ai.search("Tell me about AI", stream=True):
            ...     print(chunk, end="", flush=True)
        """
        if not stream:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                result = loop.run_until_complete(
                    self._async_search(prompt, False, raw, enable_research)
                )
                return cast(
                    Union[SearchResponse, Generator[Union[Dict[str, str], SearchResponse], None, None]],
                    result,
                )
            finally:
                loop.close()

        buffer = ""

        def sync_generator():
            nonlocal buffer
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                async_gen = loop.run_until_complete(
                    self._async_search(prompt, True, raw, enable_research)
                )
                if hasattr(async_gen, "__anext__"):
                    async_iterator = cast(AsyncIterator, async_gen)
                    while True:
                        try:
                            chunk = loop.run_until_complete(async_iterator.__anext__())
                            if isinstance(chunk, dict) and "text" in chunk:
                                buffer += chunk["text"]
                            elif isinstance(chunk, SearchResponse):
                                buffer += chunk.text
                            yield chunk
                        except StopAsyncIteration:
                            break
                elif isinstance(async_gen, SearchResponse):
                    yield async_gen
            finally:
                self.last_response = {"text": buffer}
                loop.close()

        return cast(
            Union[SearchResponse, Generator[Union[Dict[str, str], SearchResponse], None, None]],
            sync_generator(),
        )

    async def _async_search(
        self,
        prompt: str,
        stream: bool = False,
        raw: bool = False,
        enable_research: bool = False,
    ) -> Union[
        SearchResponse,
        str,
        AsyncIterator[Union[Dict[str, str], SearchResponse]],
    ]:
        """
        Internal async implementation of the search method.

        Handles the actual API communication, session management, and
        response processing.

        Args:
            prompt (str): The search query.
            stream (bool): Whether to stream the response.
            raw (bool): Whether to return raw response format.
            enable_research (bool): Whether to enable deep research.

        Returns:
            Union[SearchResponse, str, AsyncIterator]: Response data.

        Raises:
            FailedToGenerateResponseError: If API communication fails.
        """
        session = AsyncSession()
        try:
            session.headers.update(self.headers)
            if self.proxies:
                session.proxies.update(self.proxies)

            # Initialize new chat session if needed
            if not self.chat_id or not self.symmetric_key:
                self.symmetric_key = self._generate_symmetric_key()

                # Build context parameters for new session
                params = self._build_context_params(enable_research)
                params["symmetric_key"] = self.symmetric_key
                params["source"] = "home"
                params["enable_research"] = "true" if enable_research else "false"

                try:
                    response = await session.get(
                        f"{self.BASE_URL}/new",
                        params=params,
                        timeout=self.timeout,
                        impersonate="chrome110",
                    )
                    response.raise_for_status()
                except Exception as e:
                    raise exceptions.FailedToGenerateResponseError(
                        f"Failed to create new chat session: {str(e)}"
                    )

                try:
                    data = response.json()
                    self.chat_id = data.get("id")
                    if not self.chat_id:
                        raise exceptions.FailedToGenerateResponseError(
                            "No conversation ID returned from API"
                        )
                except json.JSONDecodeError as e:
                    raise exceptions.FailedToGenerateResponseError(
                        f"Invalid JSON response: {str(e)}"
                    )

            # Build stream request parameters
            params = self._build_context_params(enable_research)
            params["id"] = self.chat_id
            params["query"] = prompt
            params["symmetric_key"] = self.symmetric_key
            params["enable_followups"] = "true" if enable_research else "false"

            # Set proper referer for the request
            referer = f"https://search.brave.com/ask?q={quote(prompt)}&conversation={self.chat_id}"
            headers = {"referer": referer}

            if not stream:
                # Non-streaming mode: collect entire response
                full_text = ""
                response = await session.get(
                    f"{self.BASE_URL}/stream",
                    params=params,
                    headers=headers,
                    stream=True,
                    timeout=self.timeout,
                    impersonate="chrome110",
                )
                response.raise_for_status()

                async for line in response.aiter_lines():
                    if not line:
                        continue
                    try:
                        j = json.loads(line)
                        if j.get("type") == "text_delta":
                            full_text += j.get("delta", "")
                    except json.JSONDecodeError:
                        continue

                self.last_response = SearchResponse(full_text)
                return full_text if raw else self.last_response

            # Streaming mode: yield chunks as they arrive
            async def process_stream():
                buffer = ""
                try:
                    response = await session.get(
                        f"{self.BASE_URL}/stream",
                        params=params,
                        headers=headers,
                        stream=True,
                        timeout=self.timeout,
                        impersonate="chrome110",
                    )
                    response.raise_for_status()

                    async for line in response.aiter_lines():
                        if not line:
                            continue
                        try:
                            j = json.loads(line)
                            t = j.get("type")
                            if t == "text_delta":
                                delta = j.get("delta", "")
                                buffer += delta
                                if raw:
                                    yield j
                                else:
                                    yield SearchResponse(delta)
                        except json.JSONDecodeError:
                            continue
                    self.last_response = SearchResponse(buffer)
                finally:
                    await session.close()

            return process_stream()
        except Exception:
            if not stream:
                await session.close()
            raise


if __name__ == "__main__":
    # Example usage
    ai = BraveAI()

    # Non-streaming test
    print("=== Non-streaming Response ===")
    res = ai.search("What is Python?", stream=False)
    print(f"Result: {res}\n")

    # Streaming test
    print("=== Streaming Response ===")
    res = ai.search("Explain machine learning", stream=True)
    from collections.abc import Iterable

    if isinstance(res, Iterable):
        for chunk in res:
            print(chunk, end="", flush=True)
    else:
        print(res)

    print("\n=== Done ===")
