#
#  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  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.
#

import logging
from typing import Protocol

from rag.utils.querit_conn import Querit
from rag.utils.tavily_conn import Tavily

WEB_SEARCH_PROVIDER_TAVILY = "tavily"
WEB_SEARCH_PROVIDER_QUERIT = "querit"

logger = logging.getLogger(__name__)


class WebSearchProvider(Protocol):
    def retrieve_chunks(self, question: str) -> dict[str, list]:
        """Return web results in RAGFlow's chunk and document aggregate shape."""


def _get_api_key(prompt_config: dict, field: str) -> str:
    api_key = prompt_config.get(field)
    return api_key.strip() if isinstance(api_key, str) else ""


def has_web_search_provider(prompt_config: dict | None) -> bool:
    if not prompt_config:
        return False
    provider = prompt_config.get("web_search_provider", WEB_SEARCH_PROVIDER_TAVILY)
    if provider == WEB_SEARCH_PROVIDER_TAVILY:
        return bool(_get_api_key(prompt_config, "tavily_api_key"))
    if provider == WEB_SEARCH_PROVIDER_QUERIT:
        return bool(_get_api_key(prompt_config, "querit_api_key"))
    return False


def create_web_search_provider(prompt_config: dict | None) -> WebSearchProvider | None:
    if not prompt_config:
        logger.debug("Web search provider resolution: provider=none status=disabled")
        return None

    provider = prompt_config.get("web_search_provider", WEB_SEARCH_PROVIDER_TAVILY)
    if provider not in (WEB_SEARCH_PROVIDER_TAVILY, WEB_SEARCH_PROVIDER_QUERIT):
        logger.debug("Web search provider resolution: provider=%s status=invalid", provider)
        return None
    if not has_web_search_provider(prompt_config):
        logger.debug("Web search provider resolution: provider=%s status=disabled", provider)
        return None

    logger.debug("Web search provider resolution: provider=%s status=resolved", provider)
    if provider == WEB_SEARCH_PROVIDER_QUERIT:
        return Querit(_get_api_key(prompt_config, "querit_api_key"))
    return Tavily(_get_api_key(prompt_config, "tavily_api_key"))