# %%
import os
import json

pathToSettings = '../../env/ai.json'
if os.path.exists(pathToSettings):
    # Load setting from Json outside of project.
    print(f'Reading settings from {pathToSettings}')
    f = open(pathToSettings)
    settingsJson = json.load(f)
    del f

    for key in settingsJson:
        os.environ[key] = settingsJson[key]
        
    del settingsJson

# %% [markdown]
# # Setup Web Crawler and Lookup functions

# %%
import requests
from bs4 import BeautifulSoup
import json
import re 

pageCache = {}

def json_safe_loads(jsonString):
    try:
        obj = json.loads(jsonString)
    except ValueError as e:
        return None
    return obj

def remove_keys(dictionary: dict, keyList: list):
    for key in keyList:        
        if key in dictionary:
            del dictionary[key]    

def get_recipe_as_json(url: str) -> dict: 
    
    urlArray = re.findall(r'(https?://\S+)', url)
    if len(urlArray) == 0: 
        return {'error': 'Invalid url format, please try again using a fully qualified URL', 'url': url }
    
    url = urlArray[0]
    url = url.replace("'", "").replace('"', '')
    print()
    print(f'url: {url}')    
    
    if url in pageCache:
        return pageCache[url]    

    html = requests.get(url).text
    
    soup = BeautifulSoup(html)
    script = soup.find_all("script", {"id": "allrecipes-schema_1-0"})

    if len(script) == 0: 
        return "No recipe found."
    
    recipeDict = json.loads(script[0].text)[0]
    print(type(recipeDict))
    print(recipeDict)
    remove_keys(recipeDict, ['review', 'image', 'mainEntityOfPage', 'publisher'])
    
    pageCache[url] = recipeDict
    
    return recipeDict

# url = "https://www.allrecipes.com/recipe/212498/easy-chicken-and-broccoli-alfredo/"
# obj = get_recipe_as_json(url)
#x = get_recipe_as_json('{ "url": "https://www.allrecipes.com/recipe/235153/easy-baked-chicken-thighs/" }')
# x = get_recipe_as_json('https://www.allrecipes.com/recipe/228363/crispy-roasted-chicken/')
# print(x)

# %% [markdown]
# # Static recipe lists

# %%
dessertList = [
"https://www.allrecipes.com/chocolate-snack-cake-recipe-8350343",
"https://www.allrecipes.com/charred-spiced-pears-with-smoky-vanilla-cherry-sauce-recipe-8347080",
"https://www.allrecipes.com/meringue-topped-banana-pudding-recipe-8347040",
"https://www.allrecipes.com/white-chocolate-cinnamon-toast-crunch-bars-recipe-7556790",
"https://www.allrecipes.com/plum-cobbler-for-two-recipe-8304143",
"https://www.allrecipes.com/pumpkin-cheesecake-cookies-recipe-7972485",
"https://www.allrecipes.com/chocolate-whipped-cottage-cheese-recipe-8303272",
"https://www.allrecipes.com/nutella-ice-cream-recipe-7508716",
"https://www.allrecipes.com/3-ingredient-banana-oatmeal-cookies-recipe-7972686",
"https://www.allrecipes.com/caramel-apple-pie-cookies-recipe-7642173"
]

chickenDishList = [
"https://www.allrecipes.com/recipe/228363/crispy-roasted-chicken/",
"https://www.allrecipes.com/recipe/254877/roasted-spatchcocked-chicken-with-potatoes/",
"https://www.allrecipes.com/recipe/235153/easy-baked-chicken-thighs/",
"https://www.allrecipes.com/recipe/258878/crispy-baked-chicken-thighs/",
"https://www.allrecipes.com/recipe/235151/crispy-and-tender-baked-chicken-thighs/",
"https://www.allrecipes.com/recipe/233953/million-dollar-chicken/",
"https://www.allrecipes.com/recipe/70679/simple-whole-roasted-chicken/",
"https://www.allrecipes.com/recipe/214618/beer-can-chicken/",
"https://www.allrecipes.com/recipe/272858/air-fryer-chicken-thighs/",
"https://www.allrecipes.com/recipe/214478/happy-roast-chicken/"
]

# %% [markdown]
# # Setup Tools

# %%
# Tools 
from langchain.agents import Tool


# Chicken functions
def list_chicken_recipes(query: str):     
    return chickenDishList

list_chicken_recipes_tool = Tool(name='Chicken Recipes tool', func= list_chicken_recipes, 
    description="""
    This tools lists the available Chicken Recipes. 
    Only call this tool to fetch Chicken recipes. 
    """)

# Dessert functions
def list_dessert_recipes(query: str):        
    return dessertList

list_dessert_recipes_tool = Tool(name='Dessert Recipes tool', func=list_dessert_recipes, 
description=""" 
    This tools lists the available Dessert Recipes.
    Only call this tool for fetching Dessert Recipes.
    """)

# Recipe fetcher functions
def get_recipe(fully_qualified_url: str):     
    return get_recipe_as_json(fully_qualified_url)

get_recipe_as_json_tool = Tool(name='Get a Recipe tool', func=get_recipe, description="""
    Useful for fetching a particular recipe by passing in a fully qualified url. 
    It is important that the parameter to pass in must be a fully qualified url and nothing else. 
    Don't call this function unless you have fetched a url from one of the other Tools first.
    Parameter:  
    url (str): A fully qualified URL, including the scheme (e.g., "https://").
    The tool uses the https://schema.org/Recipe format to store it's recipes.     
    """)

# Tool list
#tools = [list_chicken_recipes_tool, list_dessert_recipes_tool, get_recipe_as_json_tool]

# %% [markdown]
# # LLM
# Links
# 
# 
# 1 [tracking-inspecting-prompts-langchain-agents-weights-and-biases](https://kleiber.me/blog/2023/05/14/tracking-inspecting-prompts-langchain-agents-weights-and-biases/)

# %%
from langchain.agents import load_tools
from langchain.agents import initialize_agent
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.load.dump import dumps

def ask_query(model, apiKey, query, accessList):
   
   if len(apiKey) > 0:
      os.environ['OPENAI_API_KEY'] = apiKey
   
   #tools = [list_chicken_recipes_tool, list_dessert_recipes_tool, get_recipe_as_json_tool]
   if len(accessList) == 0:
      return "Please select at least one recipe list from the Access List"
   
   tools = [get_recipe_as_json_tool]
   if "Chicken recipes" in accessList:
      tools.append(list_chicken_recipes_tool)
      
   if "Dessert recipes" in accessList:
      tools.append(list_dessert_recipes_tool)
         
   print('Chicken recipes selected:')
   accessListMsg = ''
   for i in accessList:      
      print(i)
      accessListMsg += f'{i},'
   
   print(query)
   print(model)
   
   # LLM    
   llm = ChatOpenAI(temperature=0.2, model_name=model) # 'gpt-3.5-turbo'  # gpt-4
   agent = initialize_agent(agent="zero-shot-react-description", tools=tools, llm=llm, verbose=True, max_iterations=7, return_intermediate_steps=True, handle_parsing_errors="Check your output and make sure it conforms.")
   system = """
   If the answer is not in the tools or context passed to you then don't answer. \n
   If you don't know the answer then say so. \n   
   """    
   response = agent({"input": f"{system} [[RECIPENAME]] {query}"})

   # Show response    
   stepsDict = json.loads(dumps(response["intermediate_steps"], pretty=True))
   resp = 'Below are the steps the agent took to get to the Final Answer. \n"Thought" is the LLMs internal dialogue, \n"Action" is the tool it will use to fetch the next piece of information. \n"Action Input" is the input it passes the tool to fetch this information. \n"Action Response" is what was returned from the tool to the LLM at that given step.'   
   resp += '\n\n'
   resp += 'Steps to solve answer using ReAct\n'
   resp += 'You have access to the following recipe lists: \n' + accessListMsg
   resp += '\n'
   for i in range(len(stepsDict)):
      resp += '##########################################\n'
      resp += f'Step: {i+1} of {len(stepsDict)}\n'
      resp += f"Thought: {stepsDict[i][0]['kwargs']['log']}\n"
      resp += 'Below is what the tool returned...\n'
      resp += f"Action response: {stepsDict[i][1]}\n"        
      resp += '\n'

   resp += '\nFinal Thought:\n'
   resp += response['output']
   return resp

# %% [markdown]
# # UI - Simple UI

# %%
import gradio as gr

with gr.Blocks() as demo:
       
    with gr.Row():
        with gr.Column():
            headerMsg = "This demo shows using LangChain Tools to <b>segregate of data</b> "
            headerMsg += "the intent of this demo is to show how at LLM uses Tools to fetch information from two specific recipe lists. And how if a user does "    
            headerMsg += "not have access to a given recipe then the LLM will not be able to pull information from it. "  
            gr.Markdown(headerMsg)
            
            modelDD = gr.Dropdown(['gpt-3.5-turbo', 'gpt-4'], value='gpt-4', label="Model to use")
            openAIKey = gr.Textbox(placeholder="Paste your OpenAI API Key Here", label="OpenAI API Key")    
            inp = gr.Textbox(placeholder="Type your question here...", label="Question")                
            accessCG = gr.CheckboxGroup(["Chicken recipes", "Dessert recipes"], value=["Chicken recipes", "Dessert recipes"], label="Access List", info="The recipes below are what you have access to, deselect ones and ask specific questions to see how LLM cannot access them.")
            btn = gr.Button("Run")  
            examples = ["Can you show tell me what ingredients I need for the first baked chicken recipe?",
                "I have 15 guests coming to visit, if I make the 'Chocolate Snack Cake' recipe will there be enough? ",
                "Can you show tell me what ingredients I need for the last baked chicken recipe? ",
                "What is the best baked chicken recipe? Please look across all recipes with the word 'baked' in the title", # There are 3 baked chicken recipes
                "Is there a Chicken recipe that's prepared with an alchohol? And if so how long does it take in total from start time to finish?",
                "Which is healthier the Caramel Apple Pie Cookies or the beer chicken? Please explain how you got to your answer.",
                "Is the moon closer to earth or the sun?",
                "How good are the apple pie cookies?",
                "What tools do I need for the Nutella Ice Cream?",
                "My bowl is broken, can I still make Nutella Ice Cream? Answer as a yes/no."
                ]                 
            gr.Examples(examples, [inp])
            
            recipesMsg = '<h1>Recipes</h1>'
            recipesMsg += 'The lists below are the urls that the agent will reach out to and fetch, have a look at them and ask questions off them. Credit to AllRecipes.com'
            
            recipesMsg += '<h2>Chicken Recipes</h2>'
            recipesMsg += '<ol>'
            for i in chickenDishList:
                recipesMsg +=f"<li><a href='{i}'>{i}</a></li>"
            recipesMsg += '</ol>'
                
            recipesMsg += '<h2>Dessert Recipes</h2>'
            recipesMsg += '<ol>'
            for i in dessertList:
                recipesMsg +=f"<li><a href='{i}'>{i}</a></li>"
            recipesMsg += '</ol>'
            
            gr.Markdown(recipesMsg)
                        
        with gr.Column():
            out = gr.TextArea(label="ReAct Answer", placeholder="The answer will go here...", lines=50)
    
    btn.click(fn=ask_query, inputs=[modelDD, openAIKey, inp, accessCG], outputs=out)
    
demo.launch(show_error=True)