import os
import streamlit as st
from typing import List, Dict
from dotenv import load_dotenv
from typing_extensions import TypedDict
from langchain_core.messages import HumanMessage
from langgraph.graph import START, StateGraph, END
from langchain_google_genai import ChatGoogleGenerativeAI

# General Information Blog
st.set_page_config(
    page_title="Latest Information Blog"
)

# Secret Key
load_dotenv(override=True)
gemini_api_key = os.getenv("GEMINI_API_KEY")


# Latest Information Blog
from search_queries_generator import generate_queries
from blog_title_generator import blog_titles_agent
from blog_content_generator import blog_content_agent
from image_fomatted_blog_generator import image_formatted_blog
from image_prompts_generator import image_prompts

from company_content_generator import company_content
from marketed_blog import marketed_blog
from seo_keywords import seo_keywords
from seo_optimized_content_generator import seo_optimized_content_with_marketing
from seo_optimized_content_generator import simple_seo_optimized_content


# LLM
llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-exp", api_key = gemini_api_key) # type:ignore


# Define State
class BlogState(TypedDict):
    topic: str
    search_queries: List[str]
    blog_topics: List[str]     
    final_topic: str
    blog_content: str
    url: str #//
    contact_us_url: str #//
    company_information: str #//
    company_marketed_blog: str #//
    seo_keywords: str #//
    seo_optimized_blog: str #//
    image_formated_blog: str  
    prompts: Dict[str, str]

    


# Graph 1: 
builder1 = StateGraph(BlogState)
builder1.add_node("queries", lambda state: generate_queries(llm, state))
builder1.add_node("blog_topic", lambda state: blog_titles_agent(llm, state))

builder1.add_edge(START, "queries")
builder1.add_edge("queries", "blog_topic")
builder1.add_edge("blog_topic", END)
graph1 = builder1.compile()

# Graph 2: 
builder2 = StateGraph(BlogState)
builder2.add_node("blog_agent", lambda state: blog_content_agent(llm, state))
builder2.add_node("seo_keyword", lambda state: seo_keywords(state)) # //
builder2.add_node("simple_seo_blog", lambda state: simple_seo_optimized_content(llm, state)) # //
builder2.add_node("image_formatted_blog", lambda state: image_formatted_blog(llm, state))
builder2.add_node("image_prompts", lambda state: image_prompts(llm, state))

builder2.add_edge(START, "blog_agent")
builder2.add_edge("blog_agent", "seo_keyword")
builder2.add_edge("seo_keyword", "simple_seo_blog")
builder2.add_edge("simple_seo_blog", "image_formatted_blog")
builder2.add_edge("image_formatted_blog", "image_prompts")
builder2.add_edge("image_prompts", END)

graph2 = builder2.compile()


# Graph 3: 
builder3 = StateGraph(BlogState)

builder3.add_node("blog_agent", lambda state: blog_content_agent(llm, state))
builder3.add_node("company", lambda state: company_content(llm, state)) #//
builder3.add_node("market_blog", lambda state: marketed_blog(llm, state)) #//
builder3.add_node("seo_keyword", lambda state: seo_keywords(state)) # //
builder3.add_node("market_seo_blog", lambda state: seo_optimized_content_with_marketing(llm, state)) # //
builder3.add_node("image_formatted_blog", lambda state: image_formatted_blog(llm, state))
builder3.add_node("image_prompts", lambda state: image_prompts(llm, state))

builder3.add_edge(START, "blog_agent")
builder3.add_edge("blog_agent", "company")
builder3.add_edge("company", "market_blog")
builder3.add_edge("market_blog", "seo_keyword")
builder3.add_edge("seo_keyword", "market_seo_blog")
builder3.add_edge("market_seo_blog", "image_formatted_blog")
builder3.add_edge("image_formatted_blog", "image_prompts")
builder3.add_edge("image_prompts", END)

graph3 = builder3.compile()



st.title("Generate Latest Information Blog")

input_data = st.text_input("Enter a Category for your blog:")

if 'titles' not in st.session_state:
    st.session_state.titles = None
if 'selected_title' not in st.session_state:
    st.session_state.selected_title = None

if input_data:  # Only proceed if input_data is not empty
    if st.button("Generate Titles"):
        try:
            topic = [HumanMessage(content=input_data)]
            result = graph1.invoke({"topic": topic})
            print(result)
            st.session_state.titles = result.get('blog_topics', [])  # Avoid KeyError
        except Exception as e:
            st.error(str(e))

if st.session_state.titles:
    st.session_state.selected_title = st.selectbox("Pick the Title for your Blog", st.session_state.titles)

    # Checkbox for adding CTAs
    add_cta = st.checkbox("Include Call to Actions")

    if add_cta:
        # Input fields for website and contact URLs
        website_url = st.text_input("Enter your website URL:")
        contact_url = st.text_input("Enter your contact us URL:")

        # Generate Blog with CTAs button
        if st.button("Generate Blog with CTAs"):
            try:
                content = graph3.invoke({
                    "final_topic": st.session_state.selected_title,
                    "url": website_url,
                    "contact_us_url": contact_url
                })
                st.markdown(content.get('image_formated_blog', "No content generated."))  # Avoid KeyError
            except Exception as e:
                st.error(str(e))
    else:
        # Generate Blog button without CTAs
        if st.button("Generate Blog"):
            try:
                content = graph2.invoke({"final_topic": st.session_state.selected_title})
                st.markdown(content.get('image_formated_blog', "No content generated."))  # Avoid KeyError
            except Exception as e:
                st.error(str(e))