import os
from langchain_groq import ChatGroq
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from typing import Dict
import gradio as gr  # Import Gradio
import shutil  # Import shutil for file operations

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

# 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
chapter_agent = create_agent("chapter writer")
outline_agent = create_agent("outline creator")
editing_agent = create_agent("editor")
proofreading_agent = create_agent("proofreader")
character_agent = create_agent("character developer")

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

def generate_outline(book_title: str, agent) -> Dict[str, str]:
    """Generate an outline for the book."""
    query = f"Create an outline for a book titled '{book_title}'."
    return agent.invoke({"input": query})

def edit_content(content: str, agent) -> str:
    """Edit the provided content."""
    query = f"Edit the following content for clarity and style:\n\n{content}"
    return agent.invoke({"input": query})

def proofread_content(content: str, agent) -> str:
    """Proofread the provided content for grammar and spelling errors."""
    query = f"Proofread the following content:\n\n{content}"
    return agent.invoke({"input": query})

def generate_character_profile(character_name: str, agent) -> str:
    """Generate a character profile for the given character name."""
    query = f"Create a character profile for '{character_name}'."
    return agent.invoke({"input": query})

def save_chapter_to_file(chapter_text: str) -> str:
    """Save the generated chapter to a .txt file."""
    file_path = "generated_chapter.txt"
    with open(file_path, "w") as file:
        file.write(chapter_text)
    return file_path  # Return the file path for downloading

# Step 5: Define Gradio interface
def gradio_interface(book_title: str, chapter_title: str, chapter_synopsis: str):
    """Gradio interface for generating book content."""
    outline = generate_outline(book_title, outline_agent)
    chapter_text = generate_chapter(chapter_title, chapter_synopsis, chapter_agent)
    edited_chapter = edit_content(chapter_text, editing_agent)
    proofread_chapter = proofread_content(edited_chapter, proofreading_agent)
    
    # Save and download buttons
    save_button = gr.Button("Save Chapter")
    download_button = gr.Button("Download Chapter")

    # Define button actions
    save_button.click(fn=save_chapter_to_file, inputs=proofread_chapter, outputs="text")
    download_button.click(fn=shutil.copy, inputs="generated_chapter.txt", outputs="text")

    return outline, chapter_text, edited_chapter, proofread_chapter

# Step 6: Launch Gradio app
if __name__ == "__main__":
    # Define the Gradio interface
    iface = gr.Interface(
        fn=gradio_interface,
        inputs=[
            gr.Textbox(label="Book Title"),
            gr.Textbox(label="Chapter Title"),
            gr.Textbox(label="Chapter Synopsis")
        ],
        outputs=[
            gr.Textbox(label="Generated Outline", lines=5),
            gr.Textbox(label="Generated Chapter", lines=5),
            gr.Textbox(label="Edited Chapter", lines=5),
            gr.Textbox(label="Proofread Chapter", lines=5)
        ],
        title="Book Generator",
        description="Generate book content using AI."
    )
    iface.launch(share=True)  # Launch the Gradio app