from typing import TypedDict
from tavily import TavilyClient
from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.graph import StateGraph, START, END
import gradio as gr
from langgraph.checkpoint.memory import InMemorySaver

import os

class GraphState(TypedDict):
    query: str
    search_results1: str
    search_results2: str
    summary1: str
    summary2: str

tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
def tavily_search(state: GraphState):
    response = tavily_client.search(
        query=state["query"],
        max_results=2
    )
    return {
        "search_results1": response["results"][0]['content'][:500],
        "search_results2": response["results"][1]['content'][:500]
    }

llm = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    google_api_key=os.getenv("GOOGLE_API_KEY"),
)

def summarize1(state: GraphState):
    prompt = f"""
    Simply summarize following:

    {state['search_results1']}
    """

    response = llm.invoke(prompt)

    return {"summary1": response.content}

def summarize2(state: GraphState):
    prompt = f"""
    Simply summarize following:

    {state['search_results2']}
    """

    response = llm.invoke(prompt)

    return {"summary2": response.content}


builder = StateGraph(GraphState)

builder.add_node("search", tavily_search)
builder.add_node("summarize", summarize1)
builder.add_node("revise", summarize2)

builder.set_entry_point("search")
builder.add_edge("search", "summarize")
builder.add_edge("summarize", "revise")
builder.add_edge("revise", END)

checkpointer = InMemorySaver()
thread = {"configurable": {"thread_id": "1"}}

graph = builder.compile(
    checkpointer=checkpointer,
    interrupt_after=["summarize"]
)


def describe_topic(topic):
    result = graph.invoke(
    {"query": topic},
    config=thread
    )
    print(result["search_results1"])
    return result["summary1"]


def notgood(topic):
    result = graph.invoke(
    None,
    config=thread
    )
    print(result["search_results2"])
    return result["summary2"]

with gr.Blocks() as demo:
    gr.Markdown("""# Ask anything and press Enter""")
    
    topic_input = gr.Textbox(
        label="Ask here"
    )

    output1 = gr.Textbox(
        label="AI description"
    )

    output2 = gr.Textbox(
        label="AI revised description"
    )

    no_btn = gr.Button("Click here to revise description")

    topic_input.submit(
        describe_topic,
        inputs=topic_input,
        outputs=output1
    )

    no_btn.click(
        notgood,
        inputs=topic_input,
        outputs=output2
    )

    gr.Markdown("""
**Note:**
I'm using Google Gemini **free API tier**. 
I sacrificed some **quality** by limiting **LLM input output** to not exceed free tier.
If "error" shows up, it means free tier exceeded.
""")

    gr.Image(
        value="graph.png"
    )               

    
    

demo.launch()