import os
import time
import json
from dotenv import load_dotenv
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_community.tools import TavilySearchResults
from langgraph.graph import MessagesState
from langchain_core.prompts import PromptTemplate
from langchain_core.messages import HumanMessage, SystemMessage
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import tools_condition
from langgraph.prebuilt import ToolNode
from langgraph.graph.state import CompiledStateGraph

# Secret Key
load_dotenv(override=True)
tavily_api_key = os.getenv("TAVILY_API_KEY")


def blog_titles_agent(llm, state):
  
  queries = state["search_queries"]
  
  def tavily_search(query: str) -> str:
    """
    Performs a search on Tavily using the provided query.

    Args:
        query (str): The search query to be executed.

    Returns:
        str: The result of the search query.
    """
    tavily = TavilySearchResults(
    max_results=5,
    search_depth="advanced",
    include_answer=True,
    include_raw_content=True,
    api_key=tavily_api_key, # type: ignore
    )
    
    text = tavily.invoke({"query": f"{query}"})
    print("Tool Call")
    return text
    
    
  tool = [tavily_search]
  
  llm_with_tools = llm.bind_tools(tool)
  
  # System message
  sys_msg = SystemMessage(content="""
  Task: You are an expert Blog Title generator. Generate a blog Title on the Topic given to you by using the examples below as a guide.

  Examples:
  - Is It Time To Replace Your Boiler
  - Cold Radiators Let’s Look at the Reasons
  - Why Invest in a New Central Heating System
  - When Is the Right Time to Replace Your Boiler
  - Common Boiler Faults That Are Repaired by Gas Engineers
  - Air Source Heat Pumps vs. Traditional Boilers: Which Is Right for You
  - What Are Smart Thermostats and Are They Worth Installing

  Context:
  - Generate only 3 Blog Titles.
  - The blog Title should be engaging.
  - Always Use the attached tool tavily_search to get the latest information about the Topic.

  Constraints:
  - Dont include any punctuation mark in the Title.

  Output Format:
  - Give the Blog Titles in the JSON List of String Format.
  - Don't include any extra information.

  """)
  
  # Node
  def assistant(state: MessagesState) -> MessagesState:
    return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}
  
  
  builder: StateGraph = StateGraph(MessagesState)

  builder.add_node("assistant", assistant)
  builder.add_node("tools", ToolNode(tool))


  builder.add_edge(START, "assistant")
  builder.add_conditional_edges(
    "assistant",
    tools_condition,
  )
  builder.add_edge("tools", "assistant")
  react_graph: CompiledStateGraph = builder.compile()
 
 
  titles = []
  for q in queries:
    # time.sleep(5)
    messages = [HumanMessage(content=f"Search this query: {q}")]
    messages = react_graph.invoke({"messages": messages})
    print("\n",messages)
    t = messages['messages'][-1].content
    t = t.replace("json", "").replace("```", "").strip()
    t = json.loads(t)
    for i in t:
      titles.append(i)
     
  print(titles) 
  return {"blog_topics": titles}


def blog_titles_llm(llm, state):
  topic = state["topic"]
  blog_titles_prompt_template = PromptTemplate.from_template("""
  Task: You are an expert Blog Titles generator. Generate a blog Title on the Category given to you by using the examples below as a guide.


  Examples:
  - Is It Time To Replace Your Boiler
  - Cold Radiators Let’s Look at the Reasons
  - Why Invest in a New Central Heating System
  - When Is the Right Time to Replace Your Boiler
  - Common Boiler Faults That Are Repaired by Gas Engineers
  - Air Source Heat Pumps vs. Traditional Boilers: Which Is Right for You
  - What Are Smart Thermostats and Are They Worth Installing

  Context:
  - Generate atleast 10 Blog Titles.
  - The blog Title should be engaging.

  Constraints:
  - Dont include any punctuation mark in the Title.

  Output Format:
  - Give the Blog Titles in the JSON List of String Format.
  - Don't include any extra information.

  Category = {category}
  """)

  prompt = blog_titles_prompt_template.invoke({"category": topic})
  response = llm.invoke(prompt)
  content = response.content
  titles = content.replace("json", "").replace("```", "").strip() # type: ignore
  titles = json.loads(titles)
  print(titles) 
  return {"blog_topics": titles}