import os
import logging

from achatbot.common.http import HTTPRequest
from achatbot.common.factory import EngineClass, EngineFactory
from achatbot.common import interface
from achatbot.modules.functions.function import FunctionManager
from achatbot.common.types import (
    SearchApiArgs, Search1ApiArgs, SerperApiArgs
)
import achatbot.modules.functions.search


class SearchBaseApi(EngineClass, interface.IFunction):
    def __init__(self) -> None:
        self.requests = HTTPRequest()

    def get_tool_call(self):
        return {
            "type": "function",
            "function": {
                "name": "web_search",
                "description": "web search by query",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "web search query"
                        }
                    },
                    "required": ["query"]
                }
            }
        }

    def execute(self, session, **args):
        return self._web_search(session, **args)

    def _web_search(self, session, **args) -> str:
        pass


class SearchFuncEnvInit:
    engine: interface.IFunction | EngineClass = None

    @staticmethod
    def initSearchEngine() -> interface.IFunction | EngineClass:
        if SearchFuncEnvInit.engine is not None:
            logging.info(f"SearchFuncEnvInit.engine already initialized {SearchFuncEnvInit.engine}")
            return SearchFuncEnvInit.engine
        tag = os.getenv('FUNC_SEARCH_TAG', "search_api")
        kwargs = SearchFuncEnvInit.map_config_func[tag]()
        SearchFuncEnvInit.engine = EngineFactory.get_engine_by_tag(EngineClass, tag, **kwargs)
        logging.info(f"initSearchEngine: {tag}, {SearchFuncEnvInit.engine}")
        return SearchFuncEnvInit.engine

    @staticmethod
    def get_serper_api_args():
        return SerperApiArgs(
            gl=os.getenv("SERPER_GL", "cn"),
            hl=os.getenv("SERPER_HL", "zh-cn"),
            page=int(os.getenv("SERPER_PAGE", "1")),
            num=int(os.getenv("SERPER_NUM", "5")),
        ).__dict__

    @staticmethod
    def get_search_api_args():
        return SearchApiArgs(
            engine=os.getenv("SEARCH_ENGINE", "google"),
            gl=os.getenv("SEARCH_GL", "cn"),
            hl=os.getenv("SEARCH_HL", "zh-cn"),
            page=int(os.getenv("SERPER_PAGE", "1")),
            num=int(os.getenv("SERPER_NUM", "5")),
        ).__dict__

    @staticmethod
    def get_search1_api_args():
        return Search1ApiArgs(
            search_service=os.getenv("SEARCH1_ENGINE", "google"),
            image=bool(os.getenv("SEARCH1_IMAGE", "")),
            crawl_results=int(os.getenv("CRAWL_RESULTS", "0")),
            max_results=int(os.getenv("MAX_RESULTS", "5")),
        ).__dict__

    # TAG : config
    map_config_func = {
        'search_api': get_search_api_args,
        'search1_api': get_search1_api_args,
        'serper_api': get_serper_api_args,
    }


@FunctionManager.functions.register('web_search')
class SearchFunc:
    @staticmethod
    def get_tool_call():
        return SearchFuncEnvInit.initSearchEngine().get_tool_call()

    @staticmethod
    def execute(session, **args):
        return SearchFuncEnvInit.initSearchEngine().execute(session, **args)
