app.py

import streamlit as st
import os
import json
import pandas as pd
import time
from datetime import datetime
from utils import initialize_session_state, save_thinking_logs

# Import refactored LangGraph agents
from agents import (
    get_claude_client,
    run_chapter_creation,
    get_research_agent,
    select_personas_for_chapter,
    get_persona_contribution_legacy as get_persona_contribution,
    synthesize_chapter_legacy as synthesize_chapter,
    ChapterWorkflowState
)

# Constants for persistent storage
AUTOSAVE_INTERVAL = 60  # seconds
PERSISTENT_DIR = "persistent_data"
BOOK_DATA_FILE = os.path.join(PERSISTENT_DIR, "book_data.json")

# Create persistent directory if it doesn't exist
os.makedirs(PERSISTENT_DIR, exist_ok=True)

# Set page configuration
st.set_page_config(
    page_title="self.api Book Writer",
    page_icon="📚",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Function to save book data to persistent storage
def save_book_data():
    try:
        with open(BOOK_DATA_FILE, 'w') as f:
            json.dump(st.session_state.book_data, f, indent=2)
        st.session_state.last_save_time = time.time()
        return True
    except Exception as e:
        print(f"Error saving book data: {str(e)}")
        return False

# Function to load book data from persistent storage
def load_book_data():
    try:
        if os.path.exists(BOOK_DATA_FILE):
            with open(BOOK_DATA_FILE, 'r') as f:
                return json.load(f)
    except Exception as e:
        print(f"Error loading book data: {str(e)}")
    return None

# Auto-save checker
def check_autosave():
    current_time = time.time()
    if not hasattr(st.session_state, 'last_save_time') or (current_time - st.session_state.last_save_time) > AUTOSAVE_INTERVAL:
        save_book_data()

# Initialize session state
initialize_session_state()

# Load saved book data if available
saved_data = load_book_data()
if saved_data:
    st.session_state.book_data = saved_data

# Main UI Layout
st.title("📚 self.api Book Writer")
st.write("A multi-persona writing assistant for your book on spirituality using API metaphors")

# Main tabs for different functionality
main_tabs = st.tabs(["Book Setup", "Chapter Creation", "Review Process", "Thinking Process"])

with main_tabs[0]:  # Book Setup tab
    st.header("Book Setup")
    
    # Book metadata form
    with st.form("book_metadata_form"):
        col1, col2 = st.columns(2)
        with col1:
            book_title = st.text_input("Book Title", value=st.session_state.book_data["title"])
            book_subtitle = st.text_input("Book Subtitle", value=st.session_state.book_data["subtitle"])
        with col2:
            book_author = st.text_input("Author Name", value=st.session_state.book_data["author"])
        
        # Book description
        book_description = st.text_area("Book Description", 
                                        value=st.session_state.book_data.get("description", ""),
                                        height=150,
                                        help="Provide a general description of the book's purpose and target audience")
        
        submit_metadata = st.form_submit_button("Save Book Metadata")
        
        if submit_metadata:
            st.session_state.book_data["title"] = book_title
            st.session_state.book_data["subtitle"] = book_subtitle
            st.session_state.book_data["author"] = book_author
            st.session_state.book_data["description"] = book_description
            save_book_data()
            st.success("Book metadata saved successfully!")
    
    # Book outline creation
    st.subheader("Book Outline Generator")
    
    with st.form("outline_generator_form"):
        st.write("Generate a complete book outline based on the self.api concept")
        
        outline_prompt = st.text_area(
            "Outline Generation Prompt", 
            value=f"""Generate a detailed outline for the book '{st.session_state.book_data['title']}' with subtitle '{st.session_state.book_data['subtitle']}'. 

The book uses API (Application Programming Interface) as a metaphor for accessing different layers of consciousness, connecting tech concepts with spiritual development in an accessible way.

Include 8-12 chapters with titles and brief descriptions. For each chapter, include 4-6 key sections.

The book should progress logically from basic concepts through increasingly deeper layers of consciousness, with practical exercises throughout.""",
            height=200
        )
        
        generate_outline = st.form_submit_button("Generate Book Outline")
        
        if generate_outline:
            with st.spinner("Generating comprehensive book outline..."):
                client = get_claude_client()
                if client:
                    try:
                        response = client.messages.create(
                            model="claude-3-7-sonnet-20250219",
                            messages=[{
                                "role": "user",
                                "content": outline_prompt
                            }],
                            temperature=0.3,
                            max_tokens=2500
                        )
                        
                        outline = response.content[0].text
                        
                        # Parse the outline into a structured format for the book data
                        chapters = {}
                        current_chapter = None
                        
                        for line in outline.split('\n'):
                            line = line.strip()
                            if line.startswith('##') or line.startswith('Chapter'):  # Chapter title
                                # Extract chapter title
                                title_parts = line.split(':', 1) if ':' in line else line.split(' ', 1)
                                if len(title_parts) > 1:
                                    chapter_title = title_parts[1].strip()
                                else:
                                    chapter_title = line.strip('#').strip()
                                
                                chapter_id = f"chapter_{len(chapters) + 1}"
                                current_chapter = {
                                    "title": chapter_title,
                                    "sections": [],
                                    "content": "",
                                    "outline": ""
                                }
                                chapters[chapter_id] = current_chapter
                            elif line and current_chapter and not line.startswith('#'):
                                if not "outline" in current_chapter or not current_chapter["outline"]:
                                    current_chapter["outline"] = line
                                else:
                                    current_chapter["outline"] += f"\n{line}"
                                
                                # Try to identify sections
                                if line.startswith('*') or line.startswith('-') or (line[0].isdigit() and line[1] == '.'):
                                    section = line.strip('*').strip('-').strip().strip('.').strip()
                                    if section and section not in current_chapter["sections"]:
                                        current_chapter["sections"].append(section)
                        
                        # Update the book data
                        st.session_state.book_data["outline"] = outline
                        st.session_state.book_data["chapters"] = chapters
                        save_book_data()
                        
                        st.success("Book outline generated successfully!")
                        st.write(outline)
                    except Exception as e:
                        st.error(f"Error generating outline: {str(e)}")
                else:
                    st.error("Claude client not available. Please check your API key.")
    
    # Display current book outline if it exists
    if "outline" in st.session_state.book_data and st.session_state.book_data["outline"]:
        with st.expander("Current Book Outline", expanded=False):
            st.markdown(st.session_state.book_data["outline"])
    
    # Chapter management
    if "chapters" in st.session_state.book_data and st.session_state.book_data["chapters"]:
        st.subheader("Chapter Management")
        
        chapters = st.session_state.book_data["chapters"]
        chapter_options = {chapter_id: f"{idx+1}. {chapter['title']}" 
                          for idx, (chapter_id, chapter) in enumerate(chapters.items())}
        
        # Display existing chapters as a table
        chapter_data = []
        for chapter_id, chapter in chapters.items():
            chapter_num = list(chapters.keys()).index(chapter_id) + 1
            chapter_status = "Complete" if chapter.get("content") else "Not Started"
            if chapter_id in st.session_state.chapter_progress:
                progress = st.session_state.chapter_progress[chapter_id]
                if 0 < progress < 100:
                    chapter_status = f"In Progress ({progress}%)"
            
            chapter_data.append({
                "Chapter": f"{chapter_num}",
                "Title": chapter["title"],
                "Status": chapter_status,
                "ID": chapter_id
            })
        
        chapter_df = pd.DataFrame(chapter_data)
        if not chapter_df.empty:
            st.dataframe(chapter_df[["Chapter", "Title", "Status"]], use_container_width=True)

with main_tabs[1]:  # Chapter Creation tab - SIMPLIFIED
    st.header("Chapter Creation")
    
    if "chapters" not in st.session_state.book_data or not st.session_state.book_data["chapters"]:
        st.warning("Please generate a book outline first in the Book Setup tab.")
    else:
        chapters = st.session_state.book_data["chapters"]
        chapter_options = {chapter_id: f"{idx+1}. {chapter['title']}" 
                          for idx, (chapter_id, chapter) in enumerate(chapters.items())}
        
        # Chapter selection
        selected_chapter_id = st.selectbox(
            "Select Chapter to Work On",
            options=list(chapter_options.keys()),
            format_func=lambda x: chapter_options[x],
            index=0
        )
        
        if selected_chapter_id:
            st.session_state.current_chapter = selected_chapter_id
            selected_chapter = chapters[selected_chapter_id]
            
            # Display chapter details
            st.subheader(f"Working on: {selected_chapter['title']}")
            st.write(selected_chapter["outline"])
            
            # Chapter sections
            if selected_chapter["sections"]:
                with st.expander("Chapter Sections", expanded=True):
                    for idx, section in enumerate(selected_chapter["sections"]):
                        st.write(f"{idx+1}. {section}")
            
            # Research section
            with st.expander("Research for Chapter", expanded=False):
                research_prompt = st.text_input(
                    "Research Query",
                    placeholder=f"Enter research query for {selected_chapter['title']}..."
                )
                
                if st.button("Conduct Research"):
                    if research_prompt:
                        with st.spinner("Researching..."):
                            # Use the new LangGraph-based research
                            research_agent = get_research_agent()
                            if research_agent:
                                try:
                                    research_results = research_agent.invoke({
                                        "input": f"Research for book chapter '{selected_chapter['title']}': {research_prompt}",
                                        "chat_history": []
                                    })
                                    
                                    # Store research results
                                    if "research" not in selected_chapter:
                                        selected_chapter["research"] = []
                                    
                                    selected_chapter["research"].append({
                                        "query": research_prompt,
                                        "results": research_results["output"],
                                        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                                    })
                                    
                                    # Update the book data
                                    st.session_state.book_data["chapters"][selected_chapter_id] = selected_chapter
                                    save_book_data()
                                    save_thinking_logs()
                                    
                                    st.success("Research completed!")
                                    st.write(research_results["output"])
                                except Exception as e:
                                    st.error(f"Error conducting research: {str(e)}")
                            else:
                                st.error("Research agent not available. Please check your Tavily API key.")
                
                # Display previous research results
                if "research" in selected_chapter and selected_chapter["research"]:
                    st.subheader("Previous Research Results")
                    for idx, research in enumerate(selected_chapter["research"]):
                        with st.expander(f"Research: {research['query']}", expanded=False):
                            st.write(f"**Query:** {research['query']}")
                            st.write(f"**Timestamp:** {research['timestamp']}")
                            st.write("**Results:**")
                            st.write(research['results'])
            
            # Chapter creation options - NEW INTEGRATED WORKFLOW
            st.subheader("Chapter Creation")
            creation_options = st.radio(
                "Creation Method",
                ["Step-by-Step", "Automatic (LangGraph Workflow)"],
                horizontal=True
            )
            
            if creation_options == "Automatic (LangGraph Workflow)":
                # Integrated LangGraph workflow
                num_personas = st.number_input(
                    "Number of personas to include",
                    min_value=3,
                    max_value=7,
                    value=5
                )
                
                # Get research query if any
                research_query = None
                if st.checkbox("Include research in chapter creation"):
                    research_query = st.text_input(
                        "Research Query for Chapter Creation",
                        placeholder="Enter a research topic related to this chapter..."
                    )
                
                if st.button("🚀 Generate Complete Chapter with LangGraph"):
                    with st.spinner("Creating chapter with LangGraph workflow (this may take several minutes)..."):
                        progress_bar = st.progress(0)
                        status_placeholder = st.empty()
                        
                        try:
                            # Setup progress tracking
                            progress_stages = ["Initializing", "Research", "Persona Selection", 
                                               "Generating Contributions", "Synthesis"]
                            
                            # Update initial progress
                            progress_bar.progress(0.05)
                            status_placeholder.text(f"Stage 1/5: {progress_stages[0]}")
                            
                            # Run the LangGraph workflow
                            final_state = run_chapter_creation(
                                chapter_info=selected_chapter,
                                research_query=research_query,
                                num_personas=num_personas
                            )
                            
                            # Update the book data with the results
                            st.session_state.book_data["chapters"][selected_chapter_id] = selected_chapter
                            save_book_data()
                            save_thinking_logs()
                            
                            # Final update
                            progress_bar.progress(1.0)
                            
                            if "error" in final_state and final_state["error"]:
                                status_placeholder.text(f"Error: {final_state['error']}")
                                st.error(f"Error during chapter creation: {final_state['error']}")
                            else:
                                status_placeholder.text("Chapter created successfully!")
                                st.success("Chapter created successfully with LangGraph!")
                                st.experimental_rerun()
                        
                        except Exception as e:
                            progress_bar.progress(1.0)
                            status_placeholder.text(f"Error: {str(e)}")
                            st.error(f"Error during chapter creation: {str(e)}")
            else:
                # Step-by-Step Creation (original method with refactored components)
                col1, col2 = st.columns(2)
                
                with col1:
                    # One-click persona selection
                    if st.button("📋 Select Optimal Personas"):
                        with st.spinner("Selecting ideal personas for this chapter..."):
                            # Automatic selection using the refactored function
                            selected_personas, rationale = select_personas_for_chapter(selected_chapter, 5)
                            
                            # Store selected personas
                            selected_chapter["selected_personas"] = selected_personas
                            selected_chapter["selection_rationale"] = rationale
                            
                            # Update the book data
                            st.session_state.book_data["chapters"][selected_chapter_id] = selected_chapter
                            save_book_data()
                            save_thinking_logs()
                            
                            st.success(f"Selected {len(selected_personas)} personas for this chapter!")
                            st.experimental_rerun()
                
                with col2:
                    # One-click chapter generation
                    if "selected_personas" not in selected_chapter or not selected_chapter["selected_personas"]:
                        st.button("🚀 Generate Chapter Draft", disabled=True, help="Select personas first")
                    else:
                        if st.button("🚀 Generate Chapter Draft"):
                            with st.spinner("Creating chapter draft (this may take a few minutes)..."):
                                progress_bar = st.progress(0)
                                status_text = st.empty()
                                
                                # 1. Generate all persona contributions
                                missing_personas = []
                                for persona_id in selected_chapter["selected_personas"]:
                                    if persona_id == "meta_agent":
                                        continue  # Skip meta agent
                                    missing_personas.append(persona_id)
                                
                                if "contributions" not in selected_chapter:
                                    selected_chapter["contributions"] = {}
                                
                                for i, persona_id in enumerate(missing_personas):
                                    if persona_id in st.session_state.persona_library:
                                        persona = st.session_state.persona_library[persona_id]
                                        status_text.text(f"Generating {persona['name']} contribution...")
                                        
                                        # Get relevant research if available
                                        research_text = ""
                                        if "research" in selected_chapter and selected_chapter["research"]:
                                            for research in selected_chapter["research"]:
                                                research_text += f"Research Query: {research['query']}\nResults: {research['results']}\n\n"
                                        
                                        # Use the refactored function
                                        contribution = get_persona_contribution(persona_id, selected_chapter, research_text)
                                        
                                        if contribution:
                                            selected_chapter["contributions"][persona_id] = contribution
                                            
                                            # Update the book data and save
                                            st.session_state.book_data["chapters"][selected_chapter_id] = selected_chapter
                                            save_book_data()
                                            save_thinking_logs()
                                        
                                        # Update progress (70% of total for contributions)
                                        progress = (i + 1) / len(missing_personas) * 0.7
                                        progress_bar.progress(progress)
                                
                                # 2. Synthesize final chapter
                                status_text.text("Synthesizing complete chapter from all contributions...")
                                progress_bar.progress(0.8)
                                
                                # Use the refactored function
                                final_chapter = synthesize_chapter(selected_chapter, selected_chapter["contributions"])
                                
                                if final_chapter:
                                    selected_chapter["content"] = final_chapter
                                    
                                    # Update the book data
                                    st.session_state.book_data["chapters"][selected_chapter_id] = selected_chapter
                                    save_book_data()
                                    save_thinking_logs()
                                    
                                    progress_bar.progress(1.0)
                                    status_text.text("Chapter completed successfully!")
                                    
                                    st.success("Chapter draft created! Review it below.")
                                    st.experimental_rerun()
            
            # Display selected personas
            if "selected_personas" in selected_chapter and selected_chapter["selected_personas"]:
                st.subheader("Selected Personas")
                
                personas = []
                for persona_id in selected_chapter["selected_personas"]:
                    if persona_id in st.session_state.persona_library:
                        personas.append(st.session_state.persona_library[persona_id]["name"])
                
                st.write(", ".join(personas))
                
                with st.expander("Selection Rationale", expanded=False):
                    st.write(selected_chapter.get("selection_rationale", "No rationale available."))
            
            # Display final chapter content if available
            if "content" in selected_chapter and selected_chapter["content"]:
                st.subheader("Chapter Draft")
                
                final_content = selected_chapter["content"]
                word_count = len(final_content.split())
                target_count = "7,500-9,000"
                
                # Display word count with visual indicator
                col1, col2 = st.columns([3, 1])
                with col1:
                    st.write(f"**Word count**: {word_count:,} words (target: {target_count})")
                    # Calculate progress as percentage of minimum target (7500)
                    progress_pct = min(word_count / 7500, 1.0)
                    st.progress(progress_pct)
                with col2:
                    if word_count < 7500:
                        st.warning(f"⚠️ {7500 - word_count:,} words short")
                    elif word_count > 9000:
                        st.info(f"ℹ️ {word_count - 9000:,} words over")
                    else:
                        st.success("✅ Target reached")
                
                # Display options
                display_option = st.radio(
                    "Display Options",
                    ["Preview", "Edit", "Export"],
                    horizontal=True
                )
                
                if display_option == "Preview":
                    # Add estimated reading time 
                    reading_time = max(1, word_count // 250)  # Assuming 250 words per minute reading speed
                    st.write(f"**Estimated reading time**: {reading_time} minutes")
                    
                    # Display the content
                    st.markdown(final_content)
                elif display_option == "Edit":
                    edited_content = st.text_area(
                        "Edit Chapter Content",
                        value=final_content,
                        height=600
                    )
                    
                    # Show word count as user edits
                    edit_word_count = len(edited_content.split())
                    col1, col2 = st.columns([3, 1])
                    with col1:
                        st.write(f"**Current word count**: {edit_word_count:,} words (target: 7,500-9,000)")
                        progress_pct = min(edit_word_count / 7500, 1.0)
                        st.progress(progress_pct)
                    with col2:
                        if edit_word_count < 7500:
                            st.warning(f"⚠️ {7500 - edit_word_count:,} words short")
                        elif edit_word_count > 9000:
                            st.info(f"ℹ️ {edit_word_count - 9000:,} words over")
                        else:
                            st.success("✅ Target reached")
                    
                    if st.button("Save Edits"):
                        selected_chapter["content"] = edited_content
                        
                        # Update the book data
                        st.session_state.book_data["chapters"][selected_chapter_id] = selected_chapter
                        save_book_data()
                        
                        st.success("Chapter edits saved successfully!")
                        st.experimental_rerun()
                elif display_option == "Export":
                    st.download_button(
                        "Download Chapter as Markdown",
                        final_content,
                        file_name=f"{selected_chapter['title']}.md",
                        mime="text/markdown"
                    )
                    
                    # Also provide HTML export option
                    html_content = f"""<!DOCTYPE html>
<html>
<head>
    <title>{selected_chapter['title']} - {st.session_state.book_data['title']}</title>
    <meta charset="utf-8">
    <style>
        body {{ font-family: Arial, sans-serif; line-height: 1.6; max-width: 800px; margin: 0 auto; padding: 20px; }}
        h1, h2, h3 {{ color: #333; }}
        code {{ background-color: #f4f4f4; padding: 2px 5px; border-radius: 3px; }}
        pre {{ background-color: #f4f4f4; padding: 15px; border-radius: 5px; overflow-x: auto; }}
        blockquote {{ border-left: 4px solid #ddd; padding-left: 15px; color: #666; }}
    </style>
</head>
<body>
    <h1>{selected_chapter['title']}</h1>
"""
                    
                    # Process markdown to HTML in a safer way
                    content_html = final_content
                    content_html = content_html.replace('# ', '<h1>').replace('\n# ', '</h1>\n<h1>')
                    content_html = content_html.replace('## ', '<h2>').replace('\n## ', '</h2>\n<h2>')
                    content_html = content_html.replace('### ', '<h3>').replace('\n### ', '</h3>\n<h3>')
                    # Fix unbalanced tags by ensuring closing tags exist
                    if '<h1>' in content_html and not '</h1>' in content_html:
                        content_html += '</h1>'
                    if '<h2>' in content_html and not '</h2>' in content_html:
                        content_html += '</h2>'
                    if '<h3>' in content_html and not '</h3>' in content_html:
                        content_html += '</h3>'
                    
                    # Handle emphasis and strong formatting with proper nesting
                    # Replace pairs of ** with <strong> and </strong>
                    i = 0
                    strong_parts = []
                    while i < len(content_html):
                        if content_html[i:i+2] == '**':
                            if len(strong_parts) % 2 == 0:
                                strong_parts.append('<strong>')
                            else:
                                strong_parts.append('</strong>')
                            i += 2
                        else:
                            strong_parts.append(content_html[i])
                            i += 1
                    content_html = ''.join(strong_parts)
                    
                    # Replace pairs of * with <em> and </em>
                    i = 0
                    em_parts = []
                    while i < len(content_html):
                        if content_html[i:i+1] == '*':
                            if len(em_parts) % 2 == 0:
                                em_parts.append('<em>')
                            else:
                                em_parts.append('</em>')
                            i += 1
                        else:
                            em_parts.append(content_html[i])
                            i += 1
                    content_html = ''.join(em_parts)
                    
                    content_html = content_html.replace('\n\n', '<br><br>')
                    
                    html_content += content_html
                    html_content += """
</body>
</html>"""
                    
                    st.download_button(
                        "Download Chapter as HTML",
                        html_content,
                        file_name=f"{selected_chapter['title']}.html",
                        mime="text/html"
                    )

with main_tabs[2]:  # Review Process tab
    st.header("Book Review and Export")
    
    if "chapters" not in st.session_state.book_data or not st.session_state.book_data["chapters"]:
        st.warning("Please generate a book outline first in the Book Setup tab.")
    else:
        # Book progress overview
        st.subheader("Book Progress")
        
        chapters = st.session_state.book_data["chapters"]
        total_chapters = len(chapters)
        completed_chapters = sum(1 for chapter in chapters.values() if "content" in chapter and chapter["content"])
        progress_percentage = int((completed_chapters / total_chapters) * 100) if total_chapters > 0 else 0
        
        st.progress(progress_percentage / 100)
        st.write(f"**{completed_chapters}** out of **{total_chapters}** chapters completed ({progress_percentage}%)")
        
        # Chapter overview table
        chapter_data = []
        for chapter_id, chapter in chapters.items():
            chapter_num = list(chapters.keys()).index(chapter_id) + 1
            
            # Determine chapter status
            if "content" in chapter and chapter["content"]:
                status = "Complete"
                word_count = len(chapter.get('content', '').split())
            elif "contributions" in chapter and chapter["contributions"]:
                status = "Draft (Needs Synthesis)"
                word_count = sum(len(contrib.split()) for contrib in chapter["contributions"].values())
            elif "selected_personas" in chapter:
                status = "Planned"
                word_count = 0
            else:
                status = "Not Started"
                word_count = 0
            
            chapter_data.append({
                "Chapter": f"{chapter_num}",
                "Title": chapter["title"],
                "Status": status,
                "Word Count": f"{word_count:,}" if word_count > 0 else "0",
                "Target %": f"{min(word_count / 7500 * 100, 100):.0f}%" if word_count > 0 else "0%",
                "ID": chapter_id
            })
        
        chapter_df = pd.DataFrame(chapter_data)
        if not chapter_df.empty:
            st.dataframe(chapter_df[["Chapter", "Title", "Status", "Word Count", "Target %"]], use_container_width=True)
        
        # Book export options
        st.subheader("Export Book")
        
        export_format = st.radio(
            "Export Format",
            ["Markdown", "HTML", "JSON"],
            horizontal=True
        )
        
        if st.button("Generate Complete Book"):
            with st.spinner("Generating complete book..."):
                # Collect all completed chapters
                book_content = f"# {st.session_state.book_data['title']}\n\n"
                book_content += f"## {st.session_state.book_data['subtitle']}\n\n"
                book_content += f"By {st.session_state.book_data['author']}\n\n"
                
                # Add table of contents
                book_content += "## Table of Contents\n\n"
                for chapter_idx, (chapter_id, chapter) in enumerate(chapters.items()):
                    book_content += f"{chapter_idx + 1}. {chapter['title']}\n"
                
                book_content += "\n\n"
                
                # Add each chapter
                for chapter_idx, (chapter_id, chapter) in enumerate(chapters.items()):
                    if "content" in chapter and chapter["content"]:
                        book_content += f"# Chapter {chapter_idx + 1}: {chapter['title']}\n\n"
                        book_content += chapter["content"]
                        book_content += "\n\n---\n\n"
                
                if export_format == "Markdown":
                    st.download_button(
                        "Download Complete Book (Markdown)",
                        book_content,
                        file_name=f"{st.session_state.book_data['title']}.md",
                        mime="text/markdown"
                    )
                elif export_format == "HTML":
                    # Convert markdown to simple HTML
                    html_head = f"""<!DOCTYPE html>
<html>
<head>
    <title>{st.session_state.book_data['title']} - {st.session_state.book_data['subtitle']}</title>
    <meta charset="utf-8">
    <style>
        body {{ font-family: Arial, sans-serif; line-height: 1.6; max-width: 800px; margin: 0 auto; padding: 20px; }}
        h1, h2, h3 {{ color: #333; }}
        code {{ background-color: #f4f4f4; padding: 2px 5px; border-radius: 3px; }}
        pre {{ background-color: #f4f4f4; padding: 15px; border-radius: 5px; overflow-x: auto; }}
        blockquote {{ border-left: 4px solid #ddd; padding-left: 15px; color: #666; }}
        .chapter {{ margin-top: 50px; border-top: 1px solid #ddd; padding-top: 20px; }}
        .toc {{ margin: 30px 0; }}
        .toc a {{ text-decoration: none; color: #0066cc; }}
    </style>
</head>
<body>
    <h1>{st.session_state.book_data['title']}</h1>
    <h2>{st.session_state.book_data['subtitle']}</h2>
    <p>By {st.session_state.book_data['author']}</p>
    
    <div class="toc">
        <h2>Table of Contents</h2>
        <ol>
"""
                    
                    # Add table of contents
                    toc_content = ""
                    for chapter_idx, (chapter_id, chapter) in enumerate(chapters.items()):
                        toc_content += f'            <li><a href="#chapter-{chapter_idx + 1}">{chapter["title"]}</a></li>\n'
                    
                    toc_content += """        </ol>
    </div>
"""
                    
                    # Add each chapter
                    chapters_content = ""
                    for chapter_idx, (chapter_id, chapter) in enumerate(chapters.items()):
                        if "content" in chapter and chapter["content"]:
                            chapters_content += f'\n    <div id="chapter-{chapter_idx + 1}" class="chapter">\n'
                            chapters_content += f'        <h1>Chapter {chapter_idx + 1}: {chapter["title"]}</h1>\n'
                            
                            # Process the chapter content - improved handling for HTML tags
                            chapter_content = chapter["content"]
                            
                            # Replace headings
                            chapter_content = chapter_content.replace('# ', '<h1>').replace('\n# ', '</h1>\n<h1>')
                            chapter_content = chapter_content.replace('## ', '<h2>').replace('\n## ', '</h2>\n<h2>')
                            chapter_content = chapter_content.replace('### ', '<h3>').replace('\n### ', '</h3>\n<h3>')
                            
                            # Ensure all heading tags are properly closed
                            if chapter_content.count('<h1>') > chapter_content.count('</h1>'):
                                chapter_content += '</h1>'
                            if chapter_content.count('<h2>') > chapter_content.count('</h2>'):
                                chapter_content += '</h2>'
                            if chapter_content.count('<h3>') > chapter_content.count('</h3>'):
                                chapter_content += '</h3>'
                            
                            # Handle bold text - balanced pairs of **
                            i = 0
                            strong_parts = []
                            strong_open = False
                            while i < len(chapter_content):
                                if i+1 < len(chapter_content) and chapter_content[i:i+2] == '**':
                                    if not strong_open:
                                        strong_parts.append('<strong>')
                                        strong_open = True
                                    else:
                                        strong_parts.append('</strong>')
                                        strong_open = False
                                    i += 2
                                else:
                                    strong_parts.append(chapter_content[i])
                                    i += 1
                            
                            # Ensure all strong tags are closed
                            if strong_open:
                                strong_parts.append('</strong>')
                            
                            chapter_content = ''.join(strong_parts)
                            
                            # Handle italic text - balanced pairs of *
                            i = 0
                            em_parts = []
                            em_open = False
                            while i < len(chapter_content):
                                if chapter_content[i:i+1] == '*':
                                    if not em_open:
                                        em_parts.append('<em>')
                                        em_open = True
                                    else:
                                        em_parts.append('</em>')
                                        em_open = False
                                    i += 1
                                else:
                                    em_parts.append(chapter_content[i])
                                    i += 1
                            
                            # Ensure all em tags are closed
                            if em_open:
                                em_parts.append('</em>')
                                
                            chapter_content = ''.join(em_parts)
                            
                            # Handle paragraph breaks
                            chapter_content = chapter_content.replace('\n\n', '<br><br>')
                            
                            chapters_content += f'        {chapter_content}\n'
                            chapters_content += '    </div>\n'
                    
                    html_foot = """
</body>
</html>"""
                    
                    html_content = html_head + toc_content + chapters_content + html_foot
                    
                    st.download_button(
                        "Download Complete Book (HTML)",
                        html_content,
                        file_name=f"{st.session_state.book_data['title']}.html",
                        mime="text/html"
                    )
                elif export_format == "JSON":
                    # Export the entire book data structure
                    book_json = json.dumps(st.session_state.book_data, indent=2)
                    st.download_button(
                        "Download Book Data (JSON)",
                        book_json,
                        file_name=f"{st.session_state.book_data['title']}.json",
                        mime="application/json"
                    )