import functools
import itertools
import random
import re
import typing
from html import unescape
from urllib.parse import unquote

import msgspec
from lxml.etree import _Element  # type: ignore
from lxml.html import HTMLParser as LHTMLParser  # type: ignore
from lxml.html import document_fromstring  # type: ignore

from chitragupt.abc.browser import DuckDuckGoImage
from chitragupt.abc.browser import DuckDuckGoImageResult
from chitragupt.abc.browser import DuckDuckGoMessage
from chitragupt.abc.browser import DuckDuckGoTextSearchResult
from chitragupt.abc.browser import DuckDuckGoVideo
from chitragupt.abc.browser import DuckDuckGoVideoResult
from chitragupt.exceptions import ChitraGuptBrowserConversationLimitError
from chitragupt.exceptions import ChitraGuptBrowserNotInitializedError
from chitragupt.exceptions import ChitraGuptBrowserRatelimitError
from chitragupt.exceptions import ChitraGuptBrowserSearchError
from chitragupt.modules.browsers import BrowserBase
from chitragupt.modules.browsers import BrowserType

REGEX_STRIP_TAGS = re.compile("<.*?>")


class DuckDuckGo(BrowserBase):
    def _get_vqd(
        self,
        keywords: str,
    ) -> str:
        """Retrieve a valid vqd for the given keywords.

        Args:
            keywords: keywords for query.

        Returns:
            str: vqd string.
        """
        resp = self.session.get(
            "https://duckduckgo.com",
            params={"q": keywords},
        )
        return self._extract_vqd(
            resp.content,
            keywords,
        )

    def _extract_vqd(
        self,
        html_bytes: bytes,
        keywords: str,
    ) -> str:
        """Extract vqd from html bytes.

        Args:
            html_bytes: The HTML response to parse.
            keywords: The keywords that were used to fetch the HTML.

        Returns:
            str: The extracted vqd string.

        Raises:
            ChitraGuptBrowserSearchError: If the vqd could not be extracted.
        """
        for (
            prefix,
            prefix_len,
            suffix,
        ) in (
            (
                b'vqd="',
                5,
                b'"',
            ),
            (
                b"vqd=",
                4,
                b"&",
            ),
            (
                b"vqd='",
                5,
                b"'",
            ),
        ):
            try:
                start = html_bytes.index(prefix) + prefix_len
                end = html_bytes.index(
                    suffix,
                    start,
                )
                return html_bytes[start:end].decode()
            except ValueError:
                continue
        raise ChitraGuptBrowserSearchError(
            f"_extract_vqd() {keywords=} Could not extract vqd."
        )

    def _normalize(
        self,
        raw_html: str,
    ) -> str:
        """Strip HTML tags from the raw_html string."""
        return (
            unescape(
                REGEX_STRIP_TAGS.sub(
                    "",
                    raw_html,
                )
            )
            if raw_html
            else ""
        )

    def _normalize_url(
        self,
        url: str,
    ) -> str:
        """Unquote URL and replace spaces with '+'."""
        return (
            unquote(url).replace(
                " ",
                "+",
            )
            if url
            else ""
        )

    def _expand_proxy_tb_alias(
        self,
        proxy: typing.Optional[str],
    ) -> typing.Optional[str]:
        """Expand "tb" to a full proxy URL if applicable.

        Args:
            proxy: The proxy parameter that was passed to the search method.

        Returns:
            The expanded proxy string if the original proxy was "tb", otherwise
            the original proxy.
        """
        return "socks5://127.0.0.1:9150" if proxy == "tb" else proxy

    def _text_html(
        self,
        keywords: str,
        region: str = "wt-wt",
        timelimit: typing.Optional[str] = None,
        max_results: typing.Optional[int] = None,
        range_loop: int = 5,
    ) -> typing.List[DuckDuckGoTextSearchResult]:
        """Searches DuckDuckGo using the HTML API and returns a list of text search results.

        Args:
            keywords: The search query.
            region: The region to search in. Defaults to "wt-wt".
            timelimit: The timelimit to search within. Defaults to None.
            max_results: The maximum number of results to return. Defaults to None.
            range_loop: The number of times to loop through the search results. Defaults to 5.

        Returns:
            A list of DuckDuckGoTextSearchResult objects.
        """
        assert keywords, "keywords is mandatory"

        payload = {
            "q": keywords,
            "s": "0",
            "o": "json",
            "api": "d.js",
            "vqd": "",
            "kl": region,
            "bing_market": region,
        }
        if timelimit:
            payload["df"] = timelimit
        cache: typing.Set[str] = set()
        results: typing.List[DuckDuckGoTextSearchResult] = []

        for _ in range(range_loop):
            resp = self.session.post(
                "https://html.duckduckgo.com/html",
                data=payload,
            )
            if b"No  results." in resp.content:
                return results

            tree: typing.Any = document_fromstring(
                resp.content,
                self.parser,
            )
            elements = tree.xpath("//div[h2]")
            if not isinstance(
                elements,
                list,
            ):
                return results

            for e in elements:  # type: ignore
                if isinstance(
                    e,
                    _Element,
                ):
                    hrefxpath: typing.Any = e.xpath("./a/@href")  # type: ignore
                    href = (
                        str(hrefxpath[0])  # type: ignore
                        if hrefxpath
                        and isinstance(
                            hrefxpath,
                            list,
                        )
                        else None
                    )
                    if (
                        href
                        and href not in cache
                        and not href.startswith(
                            (
                                "http://www.google.com/search?q=",
                                "https://duckduckgo.com/y.js?ad_domain",
                            )
                        )
                    ):
                        cache.add(href)
                        titlexpath: typing.Any = e.xpath("./h2/a/text()")  # type: ignore
                        title = (
                            str(titlexpath[0])  # type: ignore
                            if titlexpath
                            and isinstance(
                                titlexpath,
                                list,
                            )
                            else ""
                        )
                        bodyxpath: typing.Any = e.xpath("./a//text()")  # type: ignore
                        body = (
                            "".join(str(x) for x in bodyxpath)  # type: ignore
                            if bodyxpath
                            and isinstance(
                                bodyxpath,
                                list,
                            )
                            else ""
                        )
                        results.append(
                            DuckDuckGoTextSearchResult(
                                title=self._normalize(title),
                                href=self._normalize_url(href),
                                body=self._normalize(body),
                            )
                        )
                        if max_results and len(results) >= max_results:
                            return results

            npx = tree.xpath('.//div[@class="nav-link"]')
            if not npx or not max_results:
                return results
            next_page = (
                npx[-1]
                if isinstance(
                    npx,
                    list,
                )
                and npx
                else None
            )  # type: ignore
            if isinstance(
                next_page,
                _Element,
            ):
                names = next_page.xpath('.//input[@type="hidden"]/@name')  # type: ignore
                values = next_page.xpath('.//input[@type="hidden"]/@value')  # type: ignore
                if isinstance(
                    names,
                    list,
                ) and isinstance(
                    values,
                    list,
                ):
                    payload = {
                        str(n): str(v)
                        for n, v in zip(
                            names,
                            values,
                        )
                    }  # type: ignore

        return results

    def _text_lite(
        self,
        keywords: str,
        region: str = "wt-wt",
        timelimit: typing.Optional[int] = None,
        max_results: typing.Optional[int] = None,
        range_loop: int = 5,
    ) -> typing.List[DuckDuckGoTextSearchResult]:
        """Lite version of DuckDuckGo search.

        Args:
            keywords: keywords for query.
            region: wt-wt, us-en, uk-en, ru-ru, etc. Defaults to "wt-wt".
            timelimit: days, weeks, months. Defaults to None.
            max_results: max number of results. If None, returns results only from the first response. Defaults to None.
            range_loop: number of times to loop over the search results. Defaults to 5.

        Returns:
            List of dictionaries with text search results.
        """
        assert keywords, "keywords is mandatory"

        payload = {
            "q": keywords,
            "s": "0",
            "o": "json",
            "api": "d.js",
            "vqd": "",
            "kl": region,
            "bing_market": region,
        }
        if timelimit:
            payload["df"] = timelimit  # type: ignore
        cache: typing.Set[str] = set()
        results: typing.List[DuckDuckGoTextSearchResult] = []
        title = ""
        for _ in range(range_loop):
            resp = self.session.post(
                "https://lite.duckduckgo.com/lite/",
                data=payload,
            )
            if b"No more results." in resp.content:
                return results

            tree = document_fromstring(
                resp.content,
                self.parser,
            )  # type: ignore
            elements = tree.xpath("//table[last()]//tr")  # type: ignore
            if not isinstance(
                elements,
                list,
            ):
                return results

            data = zip(
                itertools.cycle(
                    range(
                        1,
                        5,
                    )
                ),
                elements,
            )  # type: ignore
            for (
                i,
                e,
            ) in data:  # type: ignore
                if isinstance(
                    e,
                    _Element,
                ):
                    if i == 1:
                        hrefxpath = e.xpath(".//a//@href")  # type: ignore
                        href = (
                            str(hrefxpath[0])  # type: ignore
                            if hrefxpath
                            and isinstance(
                                hrefxpath,
                                list,
                            )
                            else None
                        )
                        if (
                            href is None
                            or href in cache
                            or href.startswith(
                                (
                                    "http://www.google.com/search?q=",
                                    "https://duckduckgo.com/y.js?ad_domain",
                                )
                            )
                        ):
                            [
                                next(
                                    data,
                                    None,
                                )  # type: ignore
                                for _ in range(3)  # type: ignore
                            ]  # skip block(i=1,2,3,4)
                        else:
                            cache.add(href)
                            titlexpath = e.xpath(".//a//text()")  # type: ignore
                            title = (
                                str(titlexpath[0])  # type: ignore
                                if titlexpath
                                and isinstance(
                                    titlexpath,
                                    list,
                                )
                                else ""
                            )
                    elif i == 2:
                        bodyxpath = e.xpath(
                            ".//td[@class='result-snippet']//text()"
                        )  # type: ignore
                        body = (
                            "".join(str(x) for x in bodyxpath).strip()  # type: ignore
                            if bodyxpath
                            and isinstance(
                                bodyxpath,
                                list,
                            )
                            else ""
                        )
                        if href:  # type: ignore
                            results.append(
                                DuckDuckGoTextSearchResult(
                                    title=self._normalize(title),
                                    href=self._normalize_url(href),
                                    body=self._normalize(body),
                                )
                            )
                            if max_results and len(results) >= max_results:
                                return results

            next_page_s = tree.xpath(  # type: ignore
                "//form[./input[contains(@value, 'ext')]]/input[@name='s']/@value"
            )
            if not next_page_s or not max_results:
                return results
            elif isinstance(
                next_page_s,
                list,
            ):
                payload["s"] = str(next_page_s[0])  # type: ignore

        return results

    @functools.cached_property
    def parser(
        self,
    ) -> LHTMLParser:
        """
        Returns a cached instance of LHTMLParser.

        The parser is configured to:
            - remove blank text
            - remove comments
            - remove processing instructions
            - not collect IDs

        The parser is used to parse HTML content and extract results.
        """
        return LHTMLParser(
            remove_blank_text=True,
            remove_comments=True,
            remove_pis=True,
            collect_ids=False,
        )

    def chat(
        self,
        keywords: str,
        model: str = "gpt-4o-mini",
        timeout: int = 30,
    ) -> str:
        """Chat with DuckDuckGo.

        Args:
            keywords: The keywords for the query.
            model: The model to use. Defaults to "gpt-4o-mini".
            timeout: The timeout in seconds. Defaults to 30.

        Returns:
            The response from DuckDuckGo.

        Raises:
            ChitraGuptBrowserNotInitializedError: If the browser is not initialized.
            ChitraGuptBrowserSearchError: If the search fails.
            ChitraGuptBrowserRatelimitError: If the search is rate limited.
            ChitraGuptBrowserConversationLimitError: If the conversation limit is exceeded.
        """
        if not self.initialized:
            raise ChitraGuptBrowserNotInitializedError(
                f"Browser not initialized. Call {self.__class__.__name__}.initialize() first."
            )
        models_deprecated = {
            "gpt-3.5": "gpt-4o-mini",
            "llama-3-70b": "llama-3.1-70b",
        }
        if model in models_deprecated:
            model = models_deprecated[model]
        models = {
            "claude-3-haiku": "claude-3-haiku-20240307",
            "gpt-4o-mini": "gpt-4o-mini",
            "llama-3.1-70b": "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
            "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1",
        }

        # Ensure vqd is retrieved before making requests
        if not self._chat_vqd:
            resp = self.session.get(
                "https://duckduckgo.com/duckchat/v1/status",
                headers={"x-vqd-accept": "1"},
            )
            self._chat_vqd = resp.headers.get(
                "x-vqd-4",
                "",
            )

        # Append user message
        self._chat_messages.append(
            {
                "role": "user",
                "content": keywords,
            }
        )
        self._chat_tokens_count += max(
            len(keywords) // 4,
            1,
        )  # Approximate token count

        json_data: typing.Dict[
            str,
            typing.Any,
        ] = {
            "model": models[model],
            "messages": self._chat_messages,
        }

        resp = self.session.post(
            "https://duckduckgo.com/duckchat/v1/chat",
            headers={"x-vqd-4": self._chat_vqd},
            json=json_data,
            timeout=timeout,
        )
        self._chat_vqd = resp.headers.get(
            "x-vqd-4",
            "",
        )
        # Extract response data
        try:
            data = ",".join(
                x
                for line in resp.text.rstrip("[DONE]LIMT_CVRSA\n").split(
                    "data:"
                )
                if (x := line.strip())
            )
            messages = msgspec.json.decode(
                "[" + data + "]",
                type=typing.List[DuckDuckGoMessage],
            )
        except Exception as e:
            raise ChitraGuptBrowserSearchError(
                f"Failed to parse chat response: {e!s}"
            ) from e
        # Process the received messages
        results: typing.List[str] = []
        for msg in messages:
            if msg.action == "error":
                err_message = msg
                if msg.status == 429:
                    raise (
                        ChitraGuptBrowserConversationLimitError(err_message)
                        if err_message.message == "ERR_CONVERSATION_LIMIT"
                        else ChitraGuptBrowserRatelimitError(
                            "ERR_CONVERSATION_LIMIT"
                        )
                    )
                raise ChitraGuptBrowserSearchError(err_message)
            elif msg.message:
                results.append(msg.message)

        result = "".join(results)

        # Append assistant response
        self._chat_messages.append(
            {
                "role": "assistant",
                "content": result,
            }
        )
        self._chat_tokens_count += max(
            len(result) // 4,
            1,
        )  # Update token count

        return result

    def images(
        self,
        keywords: str,
        region: str = "wt-wt",
        safesearch: str = "moderate",
        timelimit: typing.Optional[str] = None,
        size: typing.Optional[str] = None,
        color: typing.Optional[str] = None,
        type_image: typing.Optional[str] = None,
        layout: typing.Optional[str] = None,
        license_image: typing.Optional[str] = None,
        max_results: typing.Optional[int] = None,
        range_loop: int = 5,
    ) -> typing.List[DuckDuckGoImage]:
        """
        Search for images on DuckDuckGo.

        Args:
            keywords: The search keywords.
            region: The region code to search in. Defaults to "wt-wt".
            safesearch: The safesearch level. Can be "on", "moderate", or "off". Defaults to "moderate".
            timelimit: The time limit for images. E.g., "d" for day, "w" for week.
            size: The size of images. E.g., "small", "medium", "large".
            color: The color filter. E.g., "color", "monochrome".
            type_image: The type of images. E.g., "photo", "clipart".
            layout: The layout of images. E.g., "square", "tall".
            license_image: The license type. E.g., "creative commons".
            max_results: The maximum number of results to return.
            range_loop: The number of requests to make for paginated results. Defaults to 5.

        Returns:
            A list of DuckDuckGoImage objects representing the search results.

        Raises:
            ChitraGuptBrowserNotInitializedError: If the browser is not initialized.
        """
        if not self.initialized:
            raise ChitraGuptBrowserNotInitializedError(
                f"Browser not initialized. Call {self.__class__.__name__}.initialize() first."
            )
        assert keywords, "keywords is mandatory"

        vqd = self._get_vqd(keywords)

        safesearch_base = {
            "on": "1",
            "moderate": "1",
            "off": "-1",
        }
        timelimit = f"time:{timelimit}" if timelimit else ""
        size = f"size:{size}" if size else ""
        color = f"color:{color}" if color else ""
        type_image = f"type:{type_image}" if type_image else ""
        layout = f"layout:{layout}" if layout else ""
        license_image = f"license:{license_image}" if license_image else ""
        payload = {
            "l": region,
            "o": "json",
            "q": keywords,
            "vqd": vqd,
            "f": f"{timelimit},{size},{color},{type_image},{layout},{license_image}",
            "p": safesearch_base[safesearch.lower()],
        }

        results: typing.List[DuckDuckGoImage] = []

        for _ in range(range_loop):
            resp = self.session.get(
                "https://duckduckgo.com/i.js",
                params=payload,
            )
            image_result = msgspec.json.decode(
                resp.content,
                type=DuckDuckGoImageResult,
            )
            results = image_result.results
            if max_results and len(results) >= max_results:
                return results
            if next is None or not max_results:
                return results
            next_img = image_result.next
            payload["s"] = next_img.split("s=")[-1].split("&")[0]  # type: ignore

        return results

    def videos(
        self,
        keywords: str,
        region: str = "wt-wt",
        safesearch: str = "moderate",
        timelimit: typing.Optional[str] = None,
        resolution: typing.Optional[str] = None,
        duration: typing.Optional[str] = None,
        license_videos: typing.Optional[str] = None,
        max_results: typing.Optional[int] = None,
    ) -> typing.List[DuckDuckGoVideo]:
        """
        Search for videos on DuckDuckGo.

        Args
        ----
            keywords: The search keywords.
            region: The region code to search in. Defaults to "wt-wt".
            safesearch: The safesearch level. Can be "on", "moderate", or "off". Defaults to "moderate".
            timelimit: The time limit for videos. E.g., "d" for day, "w" for week.
            resolution: The resolution filter. E.g., "hd", "4k", "sd".
            duration: The duration filter. E.g., "short", "medium", "long".
            license_videos: The license type. E.g., "creative commons".
            max_results: The maximum number of results to return.

        Returns
        -------
            A list of DuckDuckGoVideo objects representing the search results.

        Raises
        ------
            ChitraGuptBrowserNotInitializedError: If the browser is not initialized.
        """
        if not self.initialized:
            raise ChitraGuptBrowserNotInitializedError(
                f"Browser not initialized. Call {self.__class__.__name__}.initialize() first."
            )
        assert keywords, "keywords is mandatory"

        vqd = self._get_vqd(keywords)

        safesearch_base = {
            "on": "1",
            "moderate": "1",
            "off": "-1",
        }
        filters = ",".join(
            filter(
                None,
                [
                    f"publishedAfter:{timelimit}" if timelimit else None,
                    f"videoDefinition:{resolution}" if resolution else None,
                    f"videoDuration:{duration}" if duration else None,
                    f"videoLicense:{license_videos}"
                    if license_videos
                    else None,
                ],
            )
        )

        payload = {
            "l": region,
            "o": "json",
            "q": keywords,
            "vqd": vqd,
            "f": filters,
            "p": safesearch_base[safesearch.lower()],
        }

        results: typing.List[DuckDuckGoVideo] = []

        for _ in range(8):
            resp = self.session.get(
                "https://duckduckgo.com/v.js",
                params=payload,
            )
            if resp.status_code != 200:
                break  # Stop retrying if the response is invalid

            try:
                results_data = msgspec.json.decode(
                    resp.content,
                    type=DuckDuckGoVideoResult,
                )
            except msgspec.DecodeError:
                break  # Stop if JSON decoding fails

            if not results_data or not results_data.results:
                break  # Stop if there are no results

            results.extend(results_data.results)
            if max_results and len(results) >= max_results:
                return results[
                    :max_results
                ]  # Return only the required number of results

            next_page = getattr(
                results_data,
                "next",
                None,
            )
            if not next_page:
                break  # Stop if there is no next page

            payload["s"] = next_page.split("s=")[-1].split("&")[0]

        return results

    def text(
        self,
        keywords: str,
        region: str = "wt-wt",
        timelimit: typing.Optional[str] = None,
        backend: str = "auto",
        max_results: typing.Optional[int] = None,
        range_loop: int = 5,
    ) -> typing.List[DuckDuckGoTextSearchResult]:
        """
        Search for text on DuckDuckGo.

        Args
        ----
            keywords: The search keywords.
            region: The region code to search in. Defaults to "wt-wt".
            timelimit: The time limit for search results. E.g., "d" for day, "w" for week.
            backend: The backend to use for search. Defaults to "auto".
            max_results: The maximum number of results to return.
            range_loop: The number of times to loop over the search results.

        Returns
        -------
            A list of DuckDuckGoTextSearchResult objects representing the search results.

        Raises
        ------
            ChitraGuptBrowserNotInitializedError: If the browser is not initialized.
            ChitraGuptBrowserSearchError: If there is an error while searching.
        """
        if not self.initialized:
            raise ChitraGuptBrowserNotInitializedError(
                f"Browser not initialized. Call {self.__class__.__name__}.initialize() first."
            )
        if backend in {
            "api",
            "ecosia",
        }:
            backend = "auto"
        backends = (
            [
                "html",
                "lite",
            ]
            if backend == "auto"
            else [backend]
        )
        random.shuffle(backends)

        (
            results,
            err,
        ) = (
            [],
            None,
        )
        for b in backends:
            try:
                if b == "html":
                    results = self._text_html(
                        keywords=keywords,
                        region=region,
                        timelimit=timelimit,  # type: ignore
                        max_results=max_results,
                        range_loop=range_loop,
                    )
                elif b == "lite":
                    results = self._text_lite(
                        keywords=keywords,
                        region=region,
                        timelimit=timelimit,  # type: ignore
                        max_results=max_results,
                        range_loop=range_loop,
                    )
                return results
            except Exception as ex:
                err = ex

        raise ChitraGuptBrowserSearchError(err)

    def initialize(
        self,
        headers: typing.Optional[
            typing.Dict[
                str,
                str,
            ]
        ] = None,
        proxy: typing.Optional[str] = None,
        proxies: typing.Optional[
            typing.Union[
                typing.Dict[
                    str,
                    str,
                ],
                str,
            ]
        ] = None,
        timeout: typing.Optional[float] = 10,
        verify: bool = True,
    ) -> None:
        """
        Initializes the browser with the given parameters.

        Args
        ----
            headers (typing.Optional[typing.Dict[str, str]]): Custom headers.
            proxy (typing.Optional[str]): Proxy URL.
            proxies (typing.Optional[typing.Union[typing.Dict[str, str], str]]): Proxies.
            timeout (typing.Optional[float]): Timeout in seconds.
            verify (bool): Verification of TLS certificates.

        """
        self.session.impersonate = BrowserType.chrome120  # type: ignore
        self.session.verify = verify
        self.session.headers = headers or {}  # type: ignore
        self.session.proxies = proxies or proxy  # type: ignore
        self.session.timeout = timeout  # type: ignore
        self.session.headers["Referer"] = "https://duckduckgo.com/"
        self._chat_vqd: typing.Optional[str] = ""
        self._chat_messages: typing.List[
            typing.Dict[
                str,
                str,
            ]
        ] = []
        self._chat_tokens_count = 0
        self.initialized = True
