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

st.set_page_config(
    page_title="Custom Title Blog"
)

load_dotenv(override=True)
gemini_api_key = os.getenv("GEMINI_API_KEY")
tavily_api_key = os.getenv("TAVILY_API_KEY")


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 = ChatGoogleGenerativeAI(model="gemini-2.0-flash-exp", api_key = gemini_api_key) # type:ignore

class BlogState(TypedDict):   
    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("blog_agent", lambda state: blog_content_agent(llm, state))
builder1.add_node("seo_keyword", lambda state: seo_keywords(state)) # //
builder1.add_node("simple_seo_blog", lambda state: simple_seo_optimized_content(llm, state)) # //
builder1.add_node("image_formatted_blog", lambda state: image_formatted_blog(llm, state))
builder1.add_node("image_prompts", lambda state: image_prompts(llm, state))

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

graph1 = builder1.compile()


# Graph 2: 
builder2 = StateGraph(BlogState)

builder2.add_node("blog_agent", lambda state: blog_content_agent(llm, state))
builder2.add_node("company", lambda state: company_content(llm, state)) #//
builder2.add_node("market_blog", lambda state: marketed_blog(llm, state)) #//
builder2.add_node("seo_keyword", lambda state: seo_keywords(state)) # //
builder2.add_node("market_seo_blog", lambda state: seo_optimized_content_with_marketing(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", "company")
builder2.add_edge("company", "market_blog")
builder2.add_edge("market_blog", "seo_keyword")
builder2.add_edge("seo_keyword", "market_seo_blog")
builder2.add_edge("market_seo_blog", "image_formatted_blog")
builder2.add_edge("image_formatted_blog", "image_prompts")
builder2.add_edge("image_prompts", END)

graph2 = builder2.compile()



st.title("Generate Custom Title Blog")

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

# 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 button with CTAs
    if st.button("Generate Blog with CTAs"):
        try:
            # Assuming graph.invoke accepts a dictionary with additional CTA fields
            content = graph2.invoke({
                "final_topic": input_data,
                "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 = graph1.invoke({"final_topic": input_data})
            st.markdown(content.get('image_formated_blog', "No content generated."))  # Avoid KeyError
        except Exception as e:
            st.error(str(e))