import gradio as gr
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from dotenv import load_dotenv, find_dotenv
from langchain.schema.runnable import RunnableParallel, RunnableSequence
import os

# Load env variables (ensure GOOGLE_API_KEY is set)
load_dotenv(find_dotenv())

# Initialize model
model = ChatGoogleGenerativeAI(
    model="models/gemini-1.5-flash-latest",
    temperature=0.5
)

# Prompt templates
prompt1 = PromptTemplate(
    template='Write a tweet for twitter {topic}',
    input_variables=['topic']
)
prompt2 = PromptTemplate(
    template='Write a post for Linkedin {topic}',
    input_variables=['topic']
)

parser = StrOutputParser()

# Parallel chain setup
parallel_chain = RunnableParallel({
    'tweet': RunnableSequence(prompt1, model, parser),
    'linkedin': RunnableSequence(prompt2, model, parser)
})

# Function for Gradio interface
def generate_posts(topic):
    result = parallel_chain.invoke({'topic': topic})
    
    tweet = result['tweet']
    linkedin = result['linkedin']
    
    # Save to file
    filename = f"post_{topic.replace(' ', '_')}.txt"
    with open(filename, "w", encoding="utf-8") as f:
        f.write(f"Tweet:\n{tweet}\n\nLinkedIn Post:\n{linkedin}")
    
    return tweet, linkedin, filename

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## AI Tweet and LinkedIn Post Generator")
    
    topic_input = gr.Textbox(label="Enter Topic")
    generate_btn = gr.Button("Generate Posts")
    
    tweet_output = gr.Textbox(label="Generated Tweet", lines=3)
    linkedin_output = gr.Textbox(label="Generated LinkedIn Post", lines=5)
    file_output = gr.File(label="Download Saved Post (.txt)")
    
    generate_btn.click(fn=generate_posts, inputs=[topic_input],
                       outputs=[tweet_output, linkedin_output, file_output])

# Launch the app
demo.launch()