import os
from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI

import operator
from langgraph.graph import StateGraph, START, END
from typing import Annotated, Literal
from typing_extensions import TypedDict
from langgraph.types import Send
from pydantic import BaseModel

from langchain.tools import DuckDuckGoSearchRun
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_community.tools import WikipediaQueryRun
from langchain_experimental.utilities import PythonREPL

class SmallerQuestions(BaseModel):
    questions: list[str]

class Nature(BaseModel):
    nature: Literal["research", "computation", "interpretation"]

class OverallState(TypedDict):
    initial_question: str = ""
    smaller_questions: list[str] = []
    current_question: str = ""
    context: str = ""
    key_words: str = ""
    wiki_results: str = ""
    web_results: str = ""
    nature: Literal["research", "computation", "interpretation"] = "interpretation"
    middle_answers: Annotated[list, operator.add] = []
    answer: str = ""

#Nodes

def BreakQuestion(state: OverallState):
    prompt = """Task: 
                Break the following question down into a maximum of 7 smaller questions. 
                Your response must be a list in which each item corresponds to one smaller question.
                Beware that the answer to each one of the questions you list should lead to the next question without holdbacks.
                
                Question:
                """ + state["initial_question"]
    response = llm.with_structured_output(SmallerQuestions).invoke(prompt)
    return {"smaller_questions": response.questions,
           "current_question": response.questions[0]}

def DecideNature(state: OverallState):
    prompt = """Task:
                The following question is answerable via one of the following three approaches: "computation", "research" or "interpretation".
                Computation is a question that requires a Python interpreter; research is a question that requires a web search in sites like Wikipedia; interpretation is a question that needs only reasoning to be answered.
                Decide which of those is the best approach and return only the string corresponding to the one you choose, and nothing else.
                Question: 
                """ + state["current_question"]
    response = llm.with_structured_output(Nature).invoke(prompt)
    return {"nature": response.nature}

def KeyWords(state: OverallState):
    prompt = """Task:
                Gather the keywords from the following question to perform a search. Your response must return only keywords and nothing else.
                Question:
                """ + state["current_question"]
    response = llm.invoke(prompt)
    return {"key_words": response.content}
    
def WikiSearch(state: OverallState):
    wikipedia_api = WikipediaAPIWrapper(top_k_results=1,doc_content_chars_max=1000)
    wikipedia = WikipediaQueryRun(api_wrapper=wikipedia_api)
    result = wikipedia.run({"query": state["key_words"]})
    return {"wiki_results": result}

def WebSearch(state: OverallState):
    search_tool = DuckDuckGoSearchRun()
    result = search_tool.invoke(state["key_words"])
    return {"web_results": result}

def ContextAnswer(state: OverallState):
    context = "Wiki results: " + state.get("wiki_results", "No wiki results found") + "\n\n Web results: " + state.get("web_results", "No web results found")
    return {"context": context}

def Computation(state: OverallState):
    python_repl = PythonREPL()
    prompt = """Task:
                The following question requires a computation to be answered. Create a code that solves the problem.
                Your response must be only the code that solves the problem and nothing else.
                Question:
                """ + state["current_question"]
    response = llm.invoke(prompt)
    code_result = python_repl.run(response.content)
    structured_answer = "The result of the following problem is " + code_result + "Problem: " + state["current_question"]
    return {"middle_answers": [structured_answer]}

def Reasoning(state: OverallState):
    prompt = """Task:
                Answer the following question using your reasoning. Your response must contain only the answer to the question and nothing else.
                Question:
                """ + state["current_question"]
    response = llm.invoke(prompt)
    return {"middle_answers": [response.content]}

def AdvanceToNextQuestion(state: OverallState):
    current_index = state["smaller_questions"].index(state["current_question"])
    if current_index < len(state["smaller_questions"]) - 1:
        return {"current_question": state["smaller_questions"][current_index + 1],
               "context": ""}
    return {}

def FinalAnswer(state: OverallState):
    middle_answers = "\n\n".join(state["middle_answers"])
    prompt = """Task:
                You will receive a question and some information from which you must be able to solve the question completely.
                Do your best to keep your response as close to the given information as possible. Your response should contain only the final answer to the question and nothing else.
                Question:
                """ + state["initial_question"] + "Context: " + middle_answers
    response = llm.invoke(prompt)
    return {"answer": response.content}

#Edges

def Nature(state: OverallState):
    if state["nature"] == "computation":
        return Send("computation", state)
    elif state["nature"] == "research":
        return Send("key_words", state)
    else:
        return Send("interpretation", state)

def NextQuestion(state: OverallState):
    if state["current_question"] == state["smaller_questions"][-1]:
        return "final_answer"
    else:
        return "define_nature"

        
load_dotenv()

llm = ChatGoogleGenerativeAI(
    model="gemini-1.5-pro",  # or "gemini-1.5-flash" for faster/cheaper
    temperature=0,
    max_tokens=1024,
    timeout=30,
    max_retries=2,
    google_api_key=os.getenv("GOOGLE_API_KEY")
)

builder = StateGraph(OverallState)
builder.add_node("break_question", BreakQuestion)
builder.add_node("advance", AdvanceToNextQuestion)
builder.add_node("define_nature", DecideNature)
builder.add_node("key_words", KeyWords)
builder.add_node("wiki_search", WikiSearch)
builder.add_node("web_search", WebSearch)
builder.add_node("answer", ContextAnswer)
builder.add_node("computation", Computation)
builder.add_node("interpretation", Reasoning)
builder.add_node("final_answer", FinalAnswer)
builder.add_edge(START, "break_question")
builder.add_conditional_edges("define_nature", Nature)
builder.add_edge("key_words", "wiki_search")
builder.add_edge("key_words", "web_search")
builder.add_edge("web_search", "answer")
builder.add_edge("wiki_search", "answer")
builder.add_edge("answer", "advance")
builder.add_edge("computation", "advance")
builder.add_edge("interpretation", "advance")
builder.add_conditional_edges("advance", NextQuestion)
builder.add_edge("final_answer", END)

graph = builder.compile()