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

# Load environment variables
load_dotenv(find_dotenv())

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

# Prompt templates
prompt1 = PromptTemplate(
    template='Write a joke about {topic}',
    input_variables=['topic']
)

prompt2 = PromptTemplate(
    template='Explain the joke - {text}',
    input_variables=['text']
)

# Output parser
parser = StrOutputParser()

# Create chain sequence
joke_gen_chain = RunnableSequence(prompt1, model, parser)

parallel_chain = RunnableParallel({
    'joke': RunnablePassthrough(),
    'explanation': RunnableSequence(prompt2, model, parser)
})

final_chain = RunnableSequence(joke_gen_chain, parallel_chain)

def generate_joke(topic):
    try:
        result = final_chain.invoke({'topic': topic})
        return result['joke'], result['explanation']
    except Exception as e:
        return "Error generating joke.", str(e)

with gr.Blocks() as demo:
    gr.Markdown("# AI Joke Generator with Explanation")

    with gr.Row():
        topic_input = gr.Textbox(label="Enter a Topic", placeholder="e.g., Cricket, Programming")

    submit_btn = gr.Button("Generate")
    joke_output = gr.Textbox(label="Joke")
    explanation_output = gr.Textbox(label="Explanation")

    submit_btn.click(fn=generate_joke, inputs=topic_input, outputs=[joke_output, explanation_output])

if __name__ == "__main__":
    demo.launch()