import os
from langchain_groq import ChatGroq
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from typing import Dict
import gradio as gr  # Import Gradio

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

# Step 4: Define helper functions for structured book generation
def create_book_agent(
    model_name: str = os.getenv("MODEL_NAME", "llama-3.1-8b-instant"),  # Use environment variable for model name
    temperature: float = 0.7,
    max_tokens: int = 8000,  # Set max_tokens to 8000
    **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)  # Removed token parameter
    chain = prompt_template | llm | StrOutputParser()
    return chain

def generate_chapter(title: str, agent):
    """Generate a full chapter given a title."""
    query = f"Write a detailed chapter titled '{title}'"
    try:
        # Yield chunks of text instead of returning all at once
        for chunk in agent.invoke({"input": query}):
            yield chunk  # Yield each chunk as it is generated
    except Exception as e:
        print(f"An error occurred while generating the chapter: {e}")
        yield ""

def write_book(title: str, outline: Dict[str, str]):
    """
    Generate a complete book.
    
    Args:
        title (str): The title of the book.
        outline (Dict[str, str]): A dictionary with chapter titles as keys.

    Returns:
        str: The full book as a single string.
    """
    book = f"# {title}\n\n"
    for chapter_title in outline.keys():
        book += f"## {chapter_title}\n\n"
        agent = create_book_agent()  # Create a new agent for each chapter
        for chapter_text in generate_chapter(chapter_title, agent):
            book += chapter_text  # Append each chunk to the book
    return book  # Return the complete book

# Step 6: Gradio interface
def gradio_interface(api_key: str = ""):
    """Create a Gradio interface for book generation."""
    with gr.Blocks() as demo:
        gr.Markdown("## Book Generator")
        gr.Markdown("This application was created by iLL-Ai AaronAllton and a team of Groq agents that write books.")  # Updated note
        user_input = gr.Textbox(label="Enter Book Title, Number of Chapters, and Type of Book (e.g., 'My Book, 10, Novel')", placeholder="Title, Chapters, Type")  # Combined input
        generate_button = gr.Button("Generate Outline")
        output = gr.Textbox(label="Generated Book", interactive=False)

        def generate_book_interface(user_input):
            try:
                # Ensure the input is properly formatted
                parts = user_input.split(",")
                if len(parts) < 3:
                    return "Please provide input in the format: 'Title, Number of Chapters, Type of Book'."
                
                title = parts[0].strip()  # First part is the title
                outline = parts[1].strip() if len(parts) > 1 else "1"  # Default to 1 chapter if not provided
                book_type = parts[2].strip() if len(parts) > 2 else "General"  # Default to "General" if not provided
                
                outline_dict = {}
                if not outline.strip():  # Check if outline is empty
                    return "Please provide a valid outline."  # Prompt user for valid outline
                
                # Generate outline based on user input and book type
                num_chapters = int(outline)  # Extract number of chapters
                outline_dict = {f"{book_type} Chapter {i}": None for i in range(1, num_chapters + 1)}  # Create outline dictionary with style
                
                print(f"Processed Outline: {outline_dict}")  # Debug statement
                agent = create_book_agent(api_key)  # Create agent with user-provided API key
                return write_book(title, outline_dict)  # Call the generator
            except Exception as e:
                return f"An error occurred: {e}"

        generate_button.click(generate_book_interface, inputs=user_input, outputs=output)

    demo.launch(share=True)  # Enable sharing of the Gradio app

if __name__ == "__main__":
    gradio_interface()