import requests, json
from duckduckgo_search import ddg
from llm_utils import create_chat_completion
from langchain.tools.base import BaseTool
import executor_pb2
from store import GlobalStore
import time

class HumanInputRun:
    """Tool that adds the capability to ask user for input."""

    def __init__(self, response_queue) -> None:
        """Initialize callback handler."""
        self._response_queue = response_queue
        self._store = GlobalStore()
        self.name = "Human"
        self.description = (
            "You can ask a human for guidance when you think you "
            "got stuck or you are not sure what to do next. "
            "The input should be a question for the human."
        )


    def run(self, query: str) -> str:
        """Use the Human input tool."""
        id = self._store.get_new_id()
        self._response_queue.put(executor_pb2.ExecutorResponse(type=executor_pb2.INPUT, message=query, input_id=id))
        return self._loop(id)

    def _loop(self, id: int)-> str:
        start_time = time.time()
        while time.time() - start_time < 60:
            value = self._store.get_value(id)
            if value is not None:
                return value

            time.sleep(1)

        return "not waiting for user input"


    async def arun(self, query: str) -> str:
        """Use the Human tool asynchronously."""
        raise NotImplementedError("Human tool does not support async")


class GoogleSearch:
    def __init__(self):
        self.name = "Google Search"
        self.description = "only useful for when you need to search for something on Google"

    def run(self, query:str) -> str:
        """Return the results of a google search"""
        return self.google_search(query)

    def arun(self, query:str) -> str:
        """Return the results of a google search"""
        return self.google_search(query)

    def google_search(self, query, num_results=3):
        """Return the results of a google search"""
        search_results = []
        for j in ddg(query, max_results=num_results) or []:
            search_results.append(j)

        return json.dumps(search_results, ensure_ascii=False, indent=4)


class TiDBKB:

    def __init__(self):

        self.name = "TiDB Knowledge Base"
        self.description = "only useful for when you need to answer questions about TiDB and MySQL knowledge. Useless to answer other questions"

        async def arun(self, query: str) -> str:
            """Run query through TiDB oncall bot and parse result."""
            response = self.oncall_chat_request(query)
            return response["content"]

    def run(self, query: str) -> str:
        """Run query through TiDB oncall bot and parse result."""
        response = self.oncall_chat_request(query)
        if response is None:
            return "TiDB Knowledge Base is a bit confused, nothing to tell for now"
        
        return response["content"]
    
    def oncall_chat_request(self, content):
        url = "http://chatbot-lb-1816165188.us-west-2.elb.amazonaws.com/chat"
        headers = {"Content-Type": "application/json"}
        data = [{"role": "user", "content": content}]

        try :
            response = requests.post(url, headers=headers, data=json.dumps(data), timeout=30)
        except requests.exceptions.RequestException as e:
            print(f"Request TiDB Knowledge failed with error: {e}")
            return None


        if response.status_code == 200:
            return response.json()
        else:
            print(f"Request failed with status code: {response.status_code}")
            return None
    
class OpenAIGPT:

    def __init__(self):

        self.name = "OpenAI GPT"
        self.description = "only useful if you don't find other tools that are more suitable"

    async def arun(self, query: str) -> str:
        """Run query through TiDB data lensot and parse result."""
        return create_chat_completion([{
            "role": "user",
            "content": query
        }])

    def run(self, query: str) -> str:
        """Run query through TiDB data lens and parse result."""
        # Load the API key from the environment variable
        return create_chat_completion([{
            "role": "user",
            "content": query
        }])