import re
from typing import Any, Dict, Generator, List, Optional, Union, cast

import requests

from webscout import exceptions
from webscout.AIbase import AISearch, SearchResponse
from webscout.litagent import LitAgent
from webscout.sanitize import sanitize_stream


class webpilotai(AISearch):
    """A class to interact with the webpilotai (WebPilot) AI search API.

    webpilotai provides a web-based comprehensive search SearchResponse interface that returns AI-generated
    SearchResponses with source references and related questions. It supports both streaming and
    non-streaming SearchResponses.

    Basic Usage:
        >>> from webscout import webpilotai
        >>> ai = webpilotai()
        >>> # Non-streaming example
        >>> response = ai.search("What is Python?")
        >>> print(response)
        Python is a high-level programming language...

        >>> # Streaming example
        >>> for chunk in ai.search("Tell me about AI", stream=True):
        ...     print(chunk, end="", flush=True)
        Artificial Intelligence is...

        >>> # Raw SearchResponse format
        >>> for chunk in ai.search("Hello", stream=True, raw=True):
        ...     print(chunk)
        {'text': 'Hello'}
        {'text': ' there!'}

    Args:
        timeout (int, optional): Request timeout in seconds. Defaults to 90.
        proxies (dict, optional): Proxy configuration for requests. Defaults to None.
    """

    def __init__(
        self,
        timeout: int = 90,
        proxies: Optional[dict] = None,
    ):
        """Initialize the webpilotai API client.

        Args:
            timeout (int, optional): Request timeout in seconds. Defaults to 90.
            proxies (dict, optional): Proxy configuration for requests. Defaults to None.
        """
        self.session = requests.Session()
        self.api_endpoint = "https://api.webpilotai.com/rupee/v1/search"
        self.timeout = timeout
        self.last_response = {}

        # The 'Bearer null' is part of the API's expected headers
        self.headers = {
            'Accept': 'application/json, text/plain, */*, text/event-stream',
            'Content-Type': 'application/json;charset=UTF-8',
            'Authorization': 'Bearer null',
            'Origin': 'https://www.webpilot.ai',
            'Referer': 'https://www.webpilot.ai/',
            'User-Agent': LitAgent().random(),
        }

        self.session.headers.update(self.headers)
        self.proxies = proxies

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

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

        Args:
            prompt (str): The search query or prompt to send to the API.
            stream (bool, optional): If True, yields SearchResponse chunks as they arrive.
                                   If False, returns complete SearchResponse. Defaults to False.
            raw (bool, optional): If True, returns raw SearchResponse dictionaries with 'text' key.
                                If False, returns SearchResponse objects that convert to text automatically.
                                Defaults to False.

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

        Raises:
            APIConnectionError: If the API request fails
        """
        payload = {
            "q": prompt,
            "threadId": ""  # Empty for new search
        }

        # We'll use regex-based extraction for WebPilot events. The sanitizer will:
        #  - skip lines that are not "type":"data" events using skip_regexes
        #  - extract the "content" value using a regex that handles escaped quotes
        # The regex used below for extract_regexes is: r'"content"\s*:\s*"((?:\\.|[^"\\])*)"'
        # which captures the content string while allowing escaped characters such as \".

        def for_stream():

            try:
                with self.session.post(
                    self.api_endpoint,
                    json=payload,
                    stream=True,
                    timeout=self.timeout,
                    proxies=self.proxies
                ) as response:
                    if not response.ok:
                        raise exceptions.APIConnectionError(
                            f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
                        )

                    processed_chunks = sanitize_stream(
                        data=response.iter_content(chunk_size=1024),
                        to_json=True,
                        # Extract content from parsed JSON payload similar to Monica
                        content_extractor=lambda chunk: (
                            ((chunk.get("data") or {}).get("content") if isinstance(chunk, dict) else None)
                            or ((chunk.get("data") or {}).get("text") if isinstance(chunk, dict) else None)
                            or ((chunk.get("data") or {}).get("delta", {}).get("content") if isinstance(chunk, dict) else None)
                        ),
                        skip_markers=["event:message"],
                        yield_raw_on_error=False,
                        encoding='utf-8',
                        encoding_errors='replace',
                        line_delimiter="\n",
                        raw=raw,
                        output_formatter=None if raw else lambda x: SearchResponse(x) if isinstance(x, str) else x,
                    )

                    for chunk in processed_chunks:
                        yield chunk

            except requests.exceptions.Timeout:
                raise exceptions.APIConnectionError("Request timed out")
            except requests.exceptions.RequestException as e:
                raise exceptions.APIConnectionError(f"Request failed: {e}")

        def for_non_stream():
            try:
                with self.session.post(
                    self.api_endpoint,
                    json=payload,
                    stream=False,
                    timeout=self.timeout,
                    proxies=self.proxies,
                ) as response:
                    if not response.ok:
                        raise exceptions.APIConnectionError(
                            f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
                        )

                    if raw:
                        return response.text

                    # Process full response payload using sanitize_stream similar to streaming path
                    processed_chunks = sanitize_stream(
                        data=response.content,
                        intro_value="",
                        to_json=True,
                        strip_chars=None,
                        start_marker=None,
                        end_marker=None,
                        content_extractor=lambda chunk: (
                            ((chunk.get("data") or {}).get("content") if isinstance(chunk, dict) else None)
                            or ((chunk.get("data") or {}).get("text") if isinstance(chunk, dict) else None)
                            or ((chunk.get("data") or {}).get("delta", {}).get("content") if isinstance(chunk, dict) else None)
                        ),
                        skip_markers=["event:message"],
                        yield_raw_on_error=False,
                        encoding='utf-8',
                        encoding_errors='replace',
                        buffer_size=8192,
                        output_formatter=lambda x: SearchResponse(x) if isinstance(x, str) else x,
                    )

                    full_response = ""
                    for content_chunk in processed_chunks:
                        if content_chunk is not None and isinstance(content_chunk, str):
                            full_response += content_chunk

                    formatted_response = self.format_SearchResponse(full_response)
                    self.last_response = SearchResponse(formatted_response)
                    return self.last_response

            except requests.exceptions.RequestException as e:
                raise exceptions.APIConnectionError(f"Request failed: {e}")

        if stream:
            return for_stream()
        else:
            return for_non_stream()

    @staticmethod
    def format_SearchResponse(text: str) -> str:
        """Format the SearchResponse text for better readability.

        Args:
            text (str): The raw SearchResponse text

        Returns:
            str: Formatted text with improved structure
        """
        # Clean up formatting
        # Remove excessive newlines
        clean_text = re.sub(r'\n{3,}', '\n\n', text)

        # Ensure consistent spacing around sections
        clean_text = re.sub(r'([.!?])\s*\n\s*([A-Z])', r'\1\n\n\2', clean_text)

        # Clean up any leftover HTML or markdown artifacts
        clean_text = re.sub(r'<[^>]*>', '', clean_text)

        # Remove trailing whitespace on each line
        clean_text = '\n'.join(line.rstrip() for line in clean_text.split('\n'))

        return clean_text.strip()


if __name__ == "__main__":
    from rich import print

    ai = webpilotai()
    resp = ai.search("webscout python package details", stream=True, raw=False)
    if hasattr(resp, "__iter__") and not isinstance(resp, (str, SearchResponse)):
        for chunk in resp:
            print(chunk, end="", flush=True)
    else:
        print(resp)
