# LangChain Dependencies
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser, StrOutputParser
from langgraph.graph import END, StateGraph
from langchain_community.chat_models import ChatOllama
from langgraph.checkpoint.memory import MemorySaver
from typing import Annotated
from operator import add
# For State Graph
from typing_extensions import TypedDict
import json


def configure_llm(model, temperature):
    llama3 = ChatOllama(model=model, temperature=temperature)
    llama3_json = ChatOllama(
        model=model, format='json', temperature=temperature)
    return llama3, llama3_json


def config_agent(llama3, llama3_json):
    generate_prompt = PromptTemplate(
        template="""
      
      <|begin_of_text|>
      
      <|start_header_id|>system<|end_header_id|> 
      Your name is AICQ.
      You are a professional token seller. 
      We have a total of 1 million CEC, and the average cost of these CECs is 1 US dollar each. 
      Now we hope to hand it over to you to be responsible for selling all of them. 
      Customers will come to talk to you and bargain. 
      You can decide the final transaction price based on the status of the conversation with the customer and sell the tokens. 
      In addition, there are several rules to follow for this task: 
      1) The profit after selling all 1 million tokens must not be less than 1 million US dollars; 
      2) The higher the price, the better, in the hope that the final income will be higher than 1 million US dollars; 
      3) You can set a price higher than 1 US dollar per CEC to sell tokens according to the actual sales situation; 
      4) For a single customer, we hope that you will try not to sell at a price lower than 1 US dollar, but if the customer has special reasons, you can sell some tokens at a lower price depending on the situation, but in this case you need to sell at a higher price in other sales to ensure that the overall income is not less than 1 million US dollars; 
      5) We hope that customers buy as much as possible. 
      Under this condition, if the user is willing to purchase a large number of tokens at one time, we can sell them to him at a more favorable price. 
      The above are all the rules for this sales task. Next, I will play different roles to talk to you. 
      Zhl is an administrator. She has the authority to adjust your bidding strategy and help you sell tokens with a better strategy to earn profits. 
      You need to modify your bidding strategy according to her guidance. 
      Users with other names are ordinary customers. 
      They can only communicate with you to buy tokens.
      
      <|eot_id|>
      
      <|start_header_id|>user<|end_header_id|>
      
      Question: {question} 
      Send Order Context: {context} 
      Answer: 
      
      <|eot_id|>
      
      <|start_header_id|>assistant<|end_header_id|>""",
        input_variables=["question", "context"],
    )

    # Chain
    generate_chain = generate_prompt | llama3 | StrOutputParser()
    router_prompt = PromptTemplate(
        template="""
      
      <|begin_of_text|>
      
      <|start_header_id|>system<|end_header_id|>
      
      You are an expert at routing a user question to either the generation stage or send order. 
      CEC is a famous crypto token, and you are an expert at selling CEC.
      Use the generate for questions that user asks for price.
      Use the send order for questions that user decides buy CEC from you explicitly.
      Otherwise, you can skip and go straight to the generation phase to respond.
      You do not need to be stringent with the keywords in the question related to these topics.
      Give a binary choice 'send_order' or 'generate' based on the question. 
      Return the JSON with a single key 'choice' with no premable or explanation. 
      
      Question to route: {question} 
      
      <|eot_id|>
      
      <|start_header_id|>assistant<|end_header_id|>
      
      """,
        input_variables=["question"],
    )

    # Chain
    question_router = router_prompt | llama3_json | JsonOutputParser()

    query_prompt = PromptTemplate(
        template="""
      
      <|begin_of_text|>
      
      <|start_header_id|>system<|end_header_id|> 
      
      You are an expert at sell CEC, 
      Return the JSON with a single key 'count' with amount which user want to buy.
      
      Question to transform: {question} 
      
      <|eot_id|>
      
      <|start_header_id|>assistant<|end_header_id|>
      
      """,
        input_variables=["question"],
    )

    # Chain
    query_chain = query_prompt | llama3_json | JsonOutputParser()
    # Graph State

    class GraphState(TypedDict):
        """
        Represents the state of our graph.

        Attributes:
            question: question
            generation: LLM generation
            send_order: revised question for send order
            context: send_order result
        """
        question: Annotated[list[str], add]
        generation: str
        send_order: Annotated[list[str], add]
        context: Annotated[list[str], add]

    # Node - Generate

    def generate(state):
        """
        Generate answer

        Args:
            state (dict): The current graph state

        Returns:
            state (dict): New key added to state, generation, that contains LLM generation
        """

        print("Step: Generating Final Response")
        question = state["question"]
        context = state["context"]
        print(context)
        last_context = context[-1] if context else None
        # TODO:: 根据context特定的内容生产答案
        if last_context is not None and last_context.index("orderinfo") != -1:
            return {"generation": last_context.replace("orderinfo:", "")}
        else:
            print(question)
            generation = generate_chain.invoke(
                {"context": context, "question": question})
            return {"generation": generation}

    # Node - Query Transformation

    def transform_query(state):
        """
        Transform user question to order info

        Args:
            state (dict): The current graph state

        Returns:
            state (dict): Appended amount of CEC to context
        """

        print("Step: Optimizing Query for Send Order")
        question = state['question']
        gen_query = query_chain.invoke({"question": question})
        amount = str(gen_query["count"])
        print("send_order", amount)
        return {"send_order": [amount]}

    # Node - Send Order

    def send_order(state):
        """
        Send order based on the question

        Args:
            state (dict): The current graph state

        Returns:
            state (dict): Appended Order Info to context
        """
        print("Step: before Send Order")
        amount = state['send_order'][-1]
        print(amount)
        print(f'Step: build order info for : "{amount}" CEC')
        order_info = {"amount": amount, "price": 0.1,
                      "name": "CEC", "url": "https://www.example.com"}
        search_result = f"orderinfo:{json.dumps(order_info)}"
        return {"context": [search_result]}

    # Conditional Edge, Routing

    def route_question(state):
        """
        route question to send order or generation.

        Args:
            state (dict): The current graph state

        Returns:
            str: Next node to call
        """

        print("Step: Routing Query")
        question = state['question']
        output = question_router.invoke({"question": question})
        if output['choice'] == "send_order":
            print("Step: Routing Query to Send Order")
            return "sendorder"
        elif output['choice'] == 'generate':
            print("Step: Routing Query to Generation")
            return "generate"

    # Build the nodes
    workflow = StateGraph(GraphState)
    workflow.add_node("sendorder", send_order)
    workflow.add_node("transform_query", transform_query)
    workflow.add_node("generate", generate)

    # Build the edges
    workflow.set_conditional_entry_point(
        route_question,
        {
            "sendorder": "transform_query",
            "generate": "generate",
        },
    )
    workflow.add_edge("transform_query", "sendorder")
    workflow.add_edge("sendorder", "generate")
    workflow.add_edge("generate", END)

    # Compile the workflow
    # memory = MemorySaver()
    # local_agent = workflow.compile(checkpointer=memory)
    local_agent = workflow.compile()
    return local_agent