import os
import gradio as gr
from langchain_groq import ChatGroq
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from typing import Dict


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

# Step 4: Define helper functions for structured book generation
def create_book_agent(
    model_name: str = "llama3-70b-8192",
    temperature: float = 0.7,
    max_tokens: int = 8192,
    **kwargs
) -> ChatGroq:
    """Create a LangChain agent for book writing."""
    prompt_template = ChatPromptTemplate.from_messages([
        ("system", "You are a creative writer. Write high-quality, engaging books for any genre."),
        ("human", "{input}")
    ])
    llm = ChatGroq(model=model_name, temperature=temperature, max_tokens=max_tokens, **kwargs)
    chain = prompt_template | llm | StrOutputParser()
    return chain

def generate_chapter(title: str, synopsis: str, agent) -> str:
    """Generate a full chapter given a title and synopsis."""
    query = f"Write a detailed chapter based on the following synopsis:\n\nTitle: {title}\n\nSynopsis: {synopsis}"
    try:
        return agent.invoke({"input": query})
    except Exception as e:
        print(f"An error occurred while generating the chapter: {e}")
        return ""

def write_book(agent, title: str, outline: Dict[str, str]) -> str:
    """
    Generate a complete book.
    
    Args:
        agent: The LangChain agent for generating text.
        title (str): The title of the book.
        outline (Dict[str, str]): A dictionary with chapter titles as keys and synopses as values.

    Returns:
        str: The full book as a single string.
    """
    book = f"# {title}\n\n"
    for chapter_title, chapter_synopsis in outline.items():
        book += f"## {chapter_title}\n\n"
        chapter_text = generate_chapter(chapter_title, chapter_synopsis, agent)
        book += chapter_text + "\n\n"
    return book

# Step 5: Create the agent
book_agent = create_book_agent()

# Step 6: Define the Gradio interface
def generate_book_ui(title: str, outline: str) -> str:
    """Gradio UI function to generate a book."""
    # Convert outline string to dictionary
    outline_dict = {}
    for line in outline.splitlines():
        if line.strip():  # Ignore empty lines
            chapter_title, chapter_synopsis = line.split(":", 1)
            outline_dict[chapter_title.strip()] = chapter_synopsis.strip()
    
    # Generate the full book
    full_book = write_book(book_agent, title, outline_dict)
    
    # Save the book to a text file
    with open("generated_book.txt", "w") as f:
        f.write(full_book)
    
    return full_book, "generated_book.txt"

# Gradio UI setup
with gr.Blocks() as demo:
    gr.Markdown("## Book Generator")
    title_input = gr.Textbox(label="Book Title", placeholder="Enter the title of your book")
    outline_input = gr.Textbox(label="Book Outline", placeholder="Enter chapter titles and synopses (one per line, format: Chapter Title: Synopsis)", lines=10)
    generate_button = gr.Button("Generate Book")
    output_text = gr.Textbox(label="Generated Book", interactive=False, lines=20)
    download_file = gr.File(label="Download Book", interactive=False)

    generate_button.click(generate_book_ui, inputs=[title_input, outline_input], outputs=[output_text, download_file])

# Launch the Gradio app
if __name__ == "__main__":
    demo.launch(share=True)