# SPDX-License-Identifier: AGPL-3.0-or-later
"""

.. sidebar:: info

   Before reporting an issue with this engine,
   please consult `API error codes`_.

Tavily_ search API (AI engine). This engine implements the REST API
(`POST /search`_) and does not make use of the `Tavily Python Wrapper`_.

From the API response, this engine generates *result items* (shown in the main
result list) and an *answer result* (shown on top of the main result list).
If the *answer* from Tavily contains an image, the *answer result* is turned
into an *infobox result*.

.. attention::

   AI queries take considerably longer to process than queries to conventional
   search engines. The ``timeout`` should therefore also be set considerably
   higher, but it is not recommended to activate AI queries by default
   (set ``disabled: true``), as otherwise all user searches will have to wait
   for the AI.

.. _Tavily: https://tavily.com/
.. _Tavily Python Wrapper: https://pypi.org/project/tavily-python/
.. _POST /search: https://docs.tavily.com/docs/rest-api/api-reference#endpoint-post-search
.. _Tavily API Credit Deduction:
   https://docs.tavily.com/docs/rest-api/api-reference#tavily-api-credit-deduction-overview
.. _Getting started: https://docs.tavily.com/docs/welcome#getting-started
.. _API error codes: https://docs.tavily.com/docs/rest-api/api-reference#error-codes

Configuration
=============

The engine has the following mandatory setting:

- :py:obj:`api_key`
- :py:obj:`topic`

Optional settings are:

- :py:obj:`days`
- :py:obj:`search_depth`
- :py:obj:`max_results`
- :py:obj:`include_answer`
- :py:obj:`include_images`
- :py:obj:`include_image_descriptions`
- :py:obj:`include_domains`
- :py:obj:`exclude_domains`

Example configuration for general search queries:

.. code:: yaml

  - name: tavily
    engine: tavily
    shortcut: tav
    categories: [general, ai]
    api_key: xxxxxxxx
    topic: general
    include_images: true
    timeout: 15
    disabled: true

Example configuration for news search:

.. code:: yaml

  - name: tavily news
    engine: tavily
    shortcut: tavnews
    categories: [news, ai]
    api_key: xxxxxxxx
    topic: news
    timeout: 15
    disabled: true


Implementation
==============

"""

from json import dumps
from datetime import datetime
from flask_babel import gettext

# about
about = {
    "website": "https://tavily.com/",
    "wikidata_id": None,
    "official_api_documentation": "https://docs.tavily.com/docs/rest-api/api-reference",
    "use_official_api": True,
    "require_api_key": True,
    "results": 'JSON',
}

search_url = "https://api.tavily.com/search"
paging = False
time_range_support = True

api_key: str = "unset"
"""Tavily API Key (`Getting started`_)."""

search_depth: str = "basic"
"""The depth of the search. It can be ``basic`` or ``advanced``. Default is
``basic`` unless specified otherwise in a given method.

- have an eye on your `Tavily API Credit Deduction`_!
"""

topic: str = ""
"""The category of the search. This will determine which of Tavily's agents
will be used for the search. Currently, only ``general`` and ``news`` are
supported."""

days: int = 3
"""The number of days back from the current date to include in the search results.
This specifies the time frame of data to be retrieved. Please note that this
feature is only available when using the ``news`` search topic. Default is 3."""

max_results: int = 5
"""The maximum number of search results to return. Default is 5."""

include_answer: bool = True
"""Include a short answer to the original query, generated by an LLM based on Tavily's
search results."""

include_images: bool = False
"""Include a list of query-related images in the response. Creates an infobox
with the first image (as far as there are any images in the response) and the answer,
if ``include_answer`` is also enabled.
"""

include_image_descriptions: bool = False
"""When ``include_images`` is set to True, this option adds descriptive text for
each image."""

include_domains: list[str] = []
"""A list of domains to specifically include in the search results. Default
is ``[]``, which includes all domains."""

exclude_domains: list[str] = []
"""A list of domains to specifically exclude from the search results. Default
is ``[]``, which doesn't exclude any domains.
"""


def request(query, params):

    data = {
        "query": query,
        "api_key": api_key,
        "search_depth": search_depth,
        "topic": topic,
        "time_range": params["time_range"],
        "max_results": max_results,
        "include_images": include_images,
        "include_domains": include_domains,
        "exclude_domains": exclude_domains,
    }

    if include_images:
        data["include_image_descriptions"] = include_image_descriptions

    if topic == "general":
        data["include_answer"] = include_answer

    elif topic == "news":
        data["days"] = days

    params["url"] = search_url
    params["method"] = "POST"
    params["headers"]["Content-type"] = "application/json"
    params["data"] = dumps(data)

    return params


def response(resp):
    results = []
    data = resp.json()

    for result in data.get("results", []):
        results.append(
            {
                "title": result['title'],
                "url": result["url"],
                "content": f"[{gettext('ai')}] {result['content']}",
                "publishedDate": _parse_date(result.get("published_date")),
            }
        )

    img_list = data.get("images")
    if img_list:
        result = {
            "infobox": f"Tavily [{gettext('ai')}]",
            "img_src": img_list[0],
        }

        content = data.get("answer")
        if isinstance(img_list[0], dict):
            result["img_src"] = img_list[0]["url"]
            img_caption = f"<i>{gettext('Image caption')}</i>: {img_list[0]['description']}"
            if not content:
                result["content"] = img_caption
            else:
                result["content"] = f"{content}<br/>{img_caption}"
        elif content:
            result["content"] = content

        results.append(result)

    elif data["answer"]:
        results.append({"answer": data["answer"]})

    return results


def _parse_date(pubDate):
    if pubDate is not None:
        try:
            return datetime.strptime(pubDate, "%a, %d %b %Y %H:%M:%S %Z")
        except (ValueError, TypeError) as e:
            logger.debug("ignore exception (publishedDate): %s", e)
    return None


def init(engine_settings: dict):
    msg = []

    val = engine_settings.get("api_key") or api_key
    if not val or val == "unset":
        msg.append("missing api_key")

    val = engine_settings.get("topic") or topic
    if val not in ["general", "news"]:
        msg.append(f"invalid topic: '{val}'")

    val = engine_settings.get("search_depth") or search_depth
    if val not in ["basic", "advanced"]:
        msg.append(f"invalid search_depth: '{val}'")

    if msg:
        raise ValueError(f"[{engine_settings['name']}] engine's settings: {' / '.join(msg)}")