import os
from dotenv import load_dotenv

from pydantic import BaseModel
from typing import List
from typing_extensions import TypedDict
from typing import Annotated,List

from langchain_groq import ChatGroq
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import ToolMessage,HumanMessage

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import AnyMessage, add_messages
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command,interrupt

load_dotenv()
os.environ["LANGCHAIN_API_KEY"] = os.getenv("LANGCHAIN_API_KEY")
os.environ["LANGCHAIN_PROJECT"] = "LangGraph Prompt Generator"
os.environ["LANGCHAIN_TRACING_V2"] = os.getenv("LANGCHAIN_TRACING_V2")
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")

# Setting up the Language Model
def get_llm():
    llm = ChatGroq(model="llama-3.3-70b-versatile", temperature=0)
    return llm

# Gathering Information from the User
def get_gather_info_llm(llm):
    class PromptInstructions(BaseModel):
        objective: str
        variables: List[str]
        constraints: List[str]
        requirements: List[str]

    llm_with_tool = llm.bind_tools([PromptInstructions])

    gather_info_prompt = """Your job is to get information from the user about what type of prompt template they want to create.
    You must get the following information from them:
    - What the objective of the prompt is
    - What variables will be passed into the prompt template
    - Any constraints for what the output should NOT do
    - Any requirements that the output MUST adhere to
    Ask user to give all the information.You must get all the above information from the user, if not clear ask the user again. 
    Do not assume any answer by yourself, get it from the user.
    After getting all the information from user, call the relevant tool.
    Don't call the tool until you get all the information from user.
    If user is not answering your question ask again"""
    gather_info_prompt_template = ChatPromptTemplate.from_messages([
        ("system", gather_info_prompt),
        MessagesPlaceholder("messages")
    ])
    gather_info_llm = gather_info_prompt_template | llm_with_tool

    return gather_info_llm

# Generating the Prompt
def get_generate_prompt_llm(llm):
    generate_prompt = """Based on the following requirements, write a good prompt template:{reqs}"""
    generate_prompt_template = ChatPromptTemplate.from_messages([
        ("system", generate_prompt)
    ])
    generate_prompt_llm = generate_prompt_template | llm

    return generate_prompt_llm

def create_graph(gather_info_llm,generate_prompt_llm):
    # State Management
    class State(TypedDict):
        messages: Annotated[List[AnyMessage], add_messages]
        is_interrupted: bool

    # Nodes
    def gather_info_node(state: State) -> State:
        state["messages"] = gather_info_llm.invoke({"messages": state["messages"]})
        
        if state["messages"].tool_calls:
            return Command(update=state, goto="tool_node")
        
        state["is_interrupted"] = True
        return Command(update=state, goto="human_in_loop_node")
    def tool_node(state: State):
        return {
            "messages": [
                ToolMessage(
                    content=str(state["messages"][-1].tool_calls[0]["args"]),
                    tool_call_id=state["messages"][-1].tool_calls[0]["id"],
                )
            ]
        }
    def generate_prompt_node(state):
        state["messages"] = generate_prompt_llm.invoke({"reqs": state["messages"][-1].content})
        return state
    def human_in_loop_node(state):
        input = interrupt("please give the requested data")
        state["is_interrupted"] = False
        state["messages"] = input
        return state

    # Building the graph
    graph_builder = StateGraph(State)

    graph_builder.add_node("gather_info_node", gather_info_node)
    graph_builder.add_node("generate_prompt_node", generate_prompt_node)
    graph_builder.add_node("tool_node", tool_node)
    graph_builder.add_node("human_in_loop_node", human_in_loop_node)

    graph_builder.add_edge(START, "gather_info_node")
    graph_builder.add_edge("human_in_loop_node", "gather_info_node")
    graph_builder.add_edge("tool_node", "generate_prompt_node")
    graph_builder.add_edge("generate_prompt_node", END)

    graph = graph_builder.compile(checkpointer=MemorySaver())

    return graph

# Implementing Graph Interaction
def invoke_graph(message, config, graph):
    human_message = HumanMessage(content=message)
    response = graph.invoke({"messages": [human_message]}, config=config)
    return response
def resume_graph(message, config, graph):
    human_message = HumanMessage(content=message)
    response = graph.invoke(Command(resume=human_message), config=config)
    return response
def display(response):
    for message in response["messages"]:
        message.pretty_print()

# Executing the Graph
llm = get_llm()
gather_info_llm= get_gather_info_llm(llm)
generate_prompt_llm= get_generate_prompt_llm(llm)
graph = create_graph(gather_info_llm,generate_prompt_llm)

config = {"configurable": {"thread_id": 1234}}

while True:
    message = input("Enter input message: ")

    graph_values = graph.get_state(config).values
    if "is_interrupted" in graph_values and graph_values["is_interrupted"]:
        response = resume_graph(message, config, graph)
    else:
        response = invoke_graph(message, config, graph)

    display(response)