import modal
import json
from openai import OpenAI
import time
from datetime import datetime
from tavily import TavilyClient
import os

tavily = TavilyClient(api_key=os.environ.get("TAVILY_API_KEY"))

client = OpenAI()
web_scraper = modal.Function.lookup("web-scraper", "extract_text")
day = datetime.now().day
month = datetime.now().strftime("%B")
year = datetime.now().year
print(f"We are the {day} of {month} {year}.")
messages = [{"role": "system", "content": "You are in charge of browsing the web for information. To reach your goal, you can use 2 differents functions: `google_search` and `scrape`. Always start with a google search except if you already know the URL of the information you are looking for. Next, you can use the `scrape` function to extract the text from any page given its url. You can repeat this process as many times as you want. When you found the information you were looking for, you can use the `return` function to send it back to the user. We are the {day} of {month} {year}."}]
functions = [
   {
      "name": "google_search",
      "description": "Search for information on Google. Returns the URL of the 10 first results.",
      "parameters": {
         "type": "object",
         "properties": {
            "query": {
               "type": "string",
               "description": "The query to search for."
            }
         },
         "required": [
            "query"
         ],
      },
   },
   {
      "name": "scrape",
      "description": "Extract text from a web page given its URL. This will return all the text on the page.",
      "parameters": {
         "type": "object",
         "properties": {
            "url": {
               "type": "string",
               "description": "The URL of the page to scrape."
            }
         },
         "required": [
            "url"
         ],
      },
   },
   {
      "name": "return",
      "description": "Return a report of the search to the user.",
      "parameters": {
         "type": "object",
         "properties": {
            "response": {
               "type": "string",
               "description": "The report to send back to the user."
            }
         },
         "required": [
            "response"
         ],
      },
   }
]

def google_search(query):
   response = tavily.search(query=query, search_depth="advanced")
   print(response)
   context = [{"url": obj["url"], "content": obj["content"]} for obj in response["results"]]
   return context

def scrape(url):
   return web_scraper.remote(url)

def browse(information: str):
   start = time.time()
   messages.append({"role": "user", "content": f"Here is the information you need to find: '{information}'"})
   while True:
      print("Calling Agent Swarm...")
      
      action = client.chat.completions.create(
         model = "gpt-4-turbo",
         messages = messages,
         temperature = 0,
         functions = functions
      ).choices[0].message
      messages.append({"role": "assistant", "content": action.content, "function_call": {"name": action.function_call.name, "arguments": action.function_call.arguments} if action.function_call is not None else None})
      if action.content is not None and action.function_call is None:
         print(action)
         print("Assistant reponded with content instead of a function call.")
         messages.append({"role": "user", "content": "Please use a function to continue."})
      elif action.function_call is not None:
         print(f"Function: {action.function_call}")
         function_name = action.function_call.name
         arguments = json.loads(action.function_call.arguments)
         if function_name == "return":
            print("Returning information to user...")
            print(arguments)
            print(f"Time elapsed: {time.time() - start} seconds")
            return {"status": "success", "response": arguments["response"]}
         if function_name not in globals() or not callable(globals()[function_name]):
            raise Exception(f"Function '{function_name}' not found.")
         function_response = eval(f"{function_name}(**{arguments})")
         print(f"Function response: {function_response}")
         messages.append({"role": "function", "content": str(function_response), "name": function_name})