import os
from langchain_groq import ChatGroq
import gradio as gr  # Ensure Gradio is imported
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from typing import Dict

# Step 1: Set the environment variable for the Groq API Key
os.environ["GROQ_API_KEY"] = "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"

# Step 2: Define a function to create agents
def create_agent(role: str, model_name: str = "llama3-70b-8192", temperature: float = 0.7) -> ChatGroq:
    """Create a LangChain agent for a specific role in book writing."""
    prompt_template = ChatPromptTemplate.from_messages([
        ("system", f"You are a {role}. Write high-quality, engaging content."),
        ("human", "{input}")
    ])
    llm = ChatGroq(model=model_name, temperature=temperature)
    chain = prompt_template | llm | StrOutputParser()
    return chain

# Step 3: Create specific agents
novel_agent = create_agent("novel writer")
guide_agent = create_agent("guide writer")

# Step 4: Define functions for generating content
def generate_novel(title: str, synopsis: str) -> str:
    """Generate a novel based on the title and synopsis."""
    query = f"Write a detailed novel based on the following synopsis:\n\nTitle: {title}\n\nSynopsis: {synopsis}"
    return novel_agent.invoke({"input": query})

def generate_guide(title: str, synopsis: str) -> str:
    """Generate a guide based on the title and synopsis."""
    query = f"Write a detailed guide based on the following synopsis:\n\nTitle: {title}\n\nSynopsis: {synopsis}"
    return guide_agent.invoke({"input": query})

# Step 5: Define Gradio interface
def gradio_interface(novel_title: str, novel_synopsis: str, guide_title: str, guide_synopsis: str):
    """Gradio interface for generating book content."""
    novel_content = generate_novel(novel_title, novel_synopsis)
    guide_content = generate_guide(guide_title, guide_synopsis)
    
    final_draft = f"Final Draft:\n\nNovel Content:\n{novel_content}\n\nGuide Content:\n{guide_content}"
    
    return novel_content, guide_content, final_draft  # Ensure all three outputs are returned

# Step 6: Launch Gradio app
if __name__ == "__main__":
    iface = gr.Interface(
        fn=gradio_interface,
        inputs=[
            gr.Textbox(label="Novel Title"),
            gr.Textbox(label="Guide Synopsis")
        ],
        outputs=[
            gr.Textbox(label="Generated Novel Content", lines=5),
            gr.Textbox(label="Generated Guide Content", lines=5),
            gr.Textbox(label="Final Draft", lines=10)  # Output for the final draft
        ],
        title="Book and Guide Generator",
        description="Generate content for novels and guides using AI."
    )
    
    # Define the function for generating the final draft
    def generate_final_draft(novel_content, guide_content):
        return f"Final Draft:\n\nNovel Content:\n{novel_content}\n\nGuide Content:\n{guide_content}"

    # Launch the Gradio app
    iface.launch(share=True)