import streamlit as st
import concurrent.futures
import asyncio
import aiohttp
from datetime import datetime
from fpdf import FPDF
import os
from dotenv import load_dotenv
from transformers import pipeline
import torch
import nltk
from nltk.tokenize import sent_tokenize, word_tokenize
from nltk.tag import pos_tag
from PIL import Image
import textwrap
import re
from typing import List, Dict, Optional
import numpy as np
from collections import defaultdict
from io import BytesIO
import matplotlib.pyplot as plt
import seaborn as sns
from faker import Faker
import wikipedia
from urllib.parse import quote

# Initialize NLTK and spaCy
nltk.download('punkt', quiet=True)
nltk.download('averaged_perceptron_tagger', quiet=True)
nltk.download('maxent_ne_chunker', quiet=True)
nltk.download('words', quiet=True)
fake = Faker()


# Load environment variables
load_dotenv()
PIXABAY_API_KEY = os.getenv('PIXABAY_API_KEY')

class EnhancedBookGenerator:
    def _enhance_content(self, content: str, facts: Optional[List[str]] = None) -> str:
        """Enhance content with facts and structure"""
        if facts is None:
            facts = []
            
        # Clean text
        content = re.sub(r'\s+', ' ', content).strip()
        
        # Process with NLTK
        sentences = sent_tokenize(content)
        
        # Add structure and facts
        structured_content = []
        current_section = []
        fact_index = 0
        
        for i, sentence in enumerate(sentences):
            current_section.append(sentence)
            
            # Add a fact every few sentences if available
            if facts and i % 5 == 0 and fact_index < len(facts):
                current_section.append(f"\nInteresting fact: {facts[fact_index]}\n")
                fact_index += 1
            
            if (i + 1) % 5 == 0:
                section_text = ' '.join(current_section)
                if self.summarizer:
                    heading = self.summarizer(section_text, max_length=10, min_length=5)[0]['summary_text']
                    structured_content.append(f"\n## {heading.title()}\n")
                structured_content.append(section_text)
                current_section = []
        
        if current_section:
            structured_content.append(' '.join(current_section))
            
        return '\n'.join(structured_content)
    def __init__(self):
        self.initialize_models()
        self.initialize_quality_metrics()
        
    def initialize_models(self):
        """Initialize NLP models with optimization"""
        try:
            # Use smaller, faster models for better performance
            self.text_generator = pipeline(
                "text-generation",
                model="distilgpt2",
                device=0 if torch.cuda.is_available() else -1
            )
            self.summarizer = pipeline(
                "summarization",
                model="facebook/bart-large-cnn",
                device=0 if torch.cuda.is_available() else -1
            )
            
        except Exception as e:
            st.error(f"Error initializing models: {str(e)}")
            self.text_generator = None
            self.summarizer = None

    def initialize_quality_metrics(self):
        """Initialize text quality metrics"""
        self.metrics = {
            'readability': defaultdict(float),
            'coherence': defaultdict(float),
            'style_match': defaultdict(float)
        }

    def process_text(self, text: str) -> str:
        """Process text using NLTK instead of spaCy"""
        sentences = sent_tokenize(text)
        words = word_tokenize(text)
        tagged = pos_tag(words)
        return ' '.join(sentences)

    async def fetch_relevant_facts(self, topic: str) -> List[str]:
        """Fetch relevant facts from Wikipedia"""
        try:
            search_results = wikipedia.search(topic, results=3)
            facts = []
            for title in search_results:
                try:
                    page = wikipedia.page(title)
                    summary = wikipedia.summary(title, sentences=2)
                    facts.append(summary)
                except:
                    continue
            return facts
        except:
            return []

    async def generate_chapter_content(self, chapter_info: Dict) -> str:
        """Generate chapter content with parallel processing"""
        try:
            prompt = self._create_chapter_prompt(chapter_info)
            
            # Parallel tasks
            tasks = [
                self._generate_raw_content(prompt, chapter_info['word_count']),
                self.fetch_relevant_facts(chapter_info['title'])
            ]
            
            # Run tasks concurrently
            results = await asyncio.gather(*tasks)
            content, facts = results
            
            # Enhance content with facts
            enhanced_content = self._enhance_content(content, facts)
            return enhanced_content
            
        except Exception as e:
            st.error(f"Error generating chapter: {str(e)}")
            return f"Error generating chapter {chapter_info['title']}"

    def _create_chapter_prompt(self, chapter_info: Dict) -> str:
        """Create detailed writing prompt with style and context"""
        base_prompt = super()._create_chapter_prompt(chapter_info)
        
        # Add user preferences if available
        if 'user_prompt' in chapter_info:
            base_prompt += f"\nConsider the following preferences: {chapter_info['user_prompt']}"
            
        return base_prompt
    
    def initialize_quality_metrics(self):
        """Initialize text quality metrics"""
        self.metrics = {
            'readability': defaultdict(float),
            'coherence': defaultdict(float),
            'style_match': defaultdict(float)
        }

    def generate_chapter_content(self, chapter_info: Dict) -> str:
        """Generate chapter content incorporating user prompt"""
        try:
            # Include user prompt in chapter generation
            base_prompt = self._create_chapter_prompt(chapter_info)
            user_prompt = st.session_state.book_settings.get('user_prompt', '')
            
            if user_prompt:
                enhanced_prompt = f"""
                User's book vision: {user_prompt}
                
                Based on this vision, write a chapter that fits the following:
                {base_prompt}
                """
            else:
                enhanced_prompt = base_prompt
                
            content = self._generate_raw_content(enhanced_prompt, chapter_info['word_count'])
            # Pass empty list for facts if none available
            enhanced_content = self._enhance_content(content, facts=[])
            return enhanced_content
            
        except Exception as e:
            st.error(f"Error generating chapter: {str(e)}")
            return f"Error generating chapter {chapter_info['title']}"

    def _create_chapter_prompt(self, chapter_info: Dict) -> str:
        """Create detailed writing prompt based on style and context"""
        style_prompts = {
            "48 Laws of Power Style": "Write an authoritative chapter about power dynamics and strategy",
            "The Prince Style": "Write a pragmatic chapter about leadership and governance",
            "Modern Business Book": "Write an insightful chapter about business strategy and success",
            "Self-Help Book": "Write an empowering chapter about personal development",
            "Novel": "Write an engaging narrative chapter with vivid descriptions"
        }
        return f"{style_prompts.get(chapter_info['style'], 'Write a chapter')} titled '{chapter_info['title']}'"

    def _generate_raw_content(self, prompt: str, word_count: int) -> str:
        """Generate content in chunks"""
        if not self.text_generator:
            return f"Sample content for {prompt}"
            
        content_parts = []
        remaining_words = min(word_count, 2000)  # Limit for stability
        
        while remaining_words > 0:
            chunk_size = min(remaining_words, 500)
            generated = self.text_generator(
                prompt,
                max_length=chunk_size,
                num_return_sequences=1,
                temperature=0.8
            )[0]['generated_text']
            content_parts.append(generated)
            remaining_words -= len(generated.split())
            
        return ' '.join(content_parts)

    def _enhance_content(self, content: str, facts: Optional[List[str]] = None) -> str:
        """Enhance content with facts and structure"""
        if facts is None:
            facts = []
            
        # Clean text
        content = re.sub(r'\s+', ' ', content).strip()
        
        # Process with NLTK
        sentences = sent_tokenize(content)
        
        # Add structure and facts
        structured_content = []
        current_section = []
        fact_index = 0
        
        for i, sentence in enumerate(sentences):
            current_section.append(sentence)
            
            # Add a fact every few sentences if available
            if facts and i % 5 == 0 and fact_index < len(facts):
                current_section.append(f"\nInteresting fact: {facts[fact_index]}\n")
                fact_index += 1
            
            if (i + 1) % 5 == 0:
                section_text = ' '.join(current_section)
                if self.summarizer:
                    heading = self.summarizer(section_text, max_length=10, min_length=5)[0]['summary_text']
                    structured_content.append(f"\n## {heading.title()}\n")
                structured_content.append(section_text)
                current_section = []
        
        if current_section:
            structured_content.append(' '.join(current_section))
            
        return '\n'.join(structured_content)

class EnhancedPDFFormatter(FPDF):
    def __init__(self, style: Dict):
        super().__init__()
        self.style = style
        self.set_auto_page_break(auto=True, margin=15)
        self.add_font('DejaVu', '', 'DejaVuSansCondensed.ttf', uni=True)
        
    def initialize_styles(self):
        """Initialize PDF styling templates"""
        self.styles = {
            "48 Laws of Power Style": {
                "title_font": ("Times", 24, "B"),
                "chapter_font": ("Times", 18, "B"),
                "body_font": ("Times", 12, ""),
                "margins": 25,
                "line_height": 1.5
            },
            "The Prince Style": {
                "title_font": ("Helvetica", 24, "B"),
                "chapter_font": ("Helvetica", 18, "B"),
                "body_font": ("Helvetica", 12, ""),
                "margins": 30,
                "line_height": 1.6
            },
            # Add more styles as needed
        }

    def chapter_title(self, title: str):
        """Add formatted chapter title"""
        self.set_font(*self.style["chapter_font"])
        self.set_fill_color(240, 240, 240)
        self.cell(0, 20, title, ln=True, fill=True)
        self.ln(10)
        
    def chapter_body(self, content: str):
        """Add formatted chapter content"""
        self.set_font(*self.style["body_font"])
        self.multi_cell(0, self.style["line_height"] * 10, content)
        self.ln()
        
    def add_footer(self):
        """Add page numbers and metadata"""
        self.set_y(-15)
        self.set_font("Arial", "I", 8)
        self.cell(0, 10, f'Page {self.page_no()}', 0, 0, 'C')

    def create_pdf(self, book_content: Dict) -> BytesIO:
        """Create formatted PDF"""
        try:
            pdf = FPDF()
            style = self.styles.get(book_content.get('style', "48 Laws of Power Style"))
            
            # Set basic formatting
            pdf.set_margins(style["margins"], style["margins"], style["margins"])
            
            # Add title page
            pdf.add_page()
            pdf.set_font(*style["title_font"])
            pdf.cell(0, 60, book_content['title'], ln=True, align='C')
            
            # Add chapters
            for chapter in book_content['chapters']:
                pdf.add_page()
                pdf.set_font(*style["chapter_font"])
                pdf.cell(0, 20, chapter['title'], ln=True)
                pdf.set_font(*style["body_font"])
                pdf.multi_cell(0, style["line_height"] * 10, chapter['content'])
            
            # Return PDF as BytesIO
            pdf_output = BytesIO()
            pdf.output(pdf_output)
            pdf_output.seek(0)
            return pdf_output
            
        except Exception as e:
            st.error(f"Error creating PDF: {str(e)}")
            return None

def initialize_session_state():
    """Initialize enhanced session state with default values"""
    if 'book_settings' not in st.session_state:
        st.session_state.book_settings = {
            'title': '',
            'template': '',
            'genre': '',
            'num_chapters': 0,
            'words_per_chapter': 0,
            'chapter_titles': [],
            'total_words': 0
        }
    if 'generated_book' not in st.session_state:
        st.session_state.generated_book = None
    if 'current_step' not in st.session_state:
        st.session_state.current_step = 'Book Settings'
    if 'user_prompt' not in st.session_state:
        st.session_state.user_prompt = ""
    if 'author_name' not in st.session_state:
        st.session_state.author_name = ""
    if 'book_generator' not in st.session_state:
        st.session_state.book_generator = EnhancedBookGenerator()

def main():
    st.set_page_config(page_title="AI Book Generator", layout="wide")
    
    # Initialize session state at the start
    initialize_session_state()

    # Book templates
    book_templates = {
        "48 Laws of Power Style": {
            "min_words": 80000,
            "max_words": 100000,
            "default_chapters": 48,
            "example_titles": ["Law 1: Never Outshine the Master", "Law 2: Never Put Too Much Trust in Friends"]
        },
        "The Prince Style": {
            "min_words": 40000,
            "max_words": 50000,
            "default_chapters": 26,
            "example_titles": ["On New Principalities", "On Military Affairs"]
        },
        "Modern Business Book": {
            "min_words": 50000,
            "max_words": 70000,
            "default_chapters": 12,
            "example_titles": ["The Art of Strategy", "Building Success"]
        },
        "Self-Help Book": {
            "min_words": 40000,
            "max_words": 60000,
            "default_chapters": 10,
            "example_titles": ["Understanding Your Potential", "Taking Action"]
        },
        "Novel": {
            "min_words": 70000,
            "max_words": 100000,
            "default_chapters": 20,
            "example_titles": ["The Beginning", "Rising Action"]
        }
    }

    # Sidebar navigation
    with st.sidebar:
        st.title("Navigation")
        selected_step = st.radio(
            "Steps",
            ["Book Settings", "Generate Content", "Preview & Download"]
        )
        st.session_state.current_step = selected_step
        
    st.title("AI Book Generator")

    # Main content area
    if st.session_state.current_step == "Book Settings":
        display_book_settings(book_templates)
    elif st.session_state.current_step == "Generate Content":
        display_generate_content()
    elif st.session_state.current_step == "Preview & Download":
        display_preview_download()

def display_book_settings(templates):
    """Display the book settings page with enhanced user prompts"""
    st.header("Book Settings")
    
    # User's creative prompt
    st.subheader("Tell us about your book idea")
    user_prompt = st.text_area(
        "Describe your book idea, themes, feelings, and key elements you want to include",
        value=st.session_state.get('user_prompt', ''),
        height=150,
        help="Be as detailed as possible. Include themes, mood, character ideas, plot elements, or any other aspects you want in your book.",
        placeholder="Example: I want to write a book about a journey of self-discovery in a dystopian world. The main themes should include hope, resilience, and the power of human connection..."
    )
    st.session_state.user_prompt = user_prompt
    
    col1, col2 = st.columns(2)
    
    with col1:
        title = st.text_input("Book Title", 
                             value=st.session_state.book_settings.get('title', ''),
                             key="book_title_input")
        
        template = st.selectbox("Writing Style", 
                              options=list(templates.keys()),
                              key="template_select")
        
        genre = st.selectbox("Genre", 
                           options=["Business", "Self-Help", "Philosophy", 
                                  "Fiction", "Non-Fiction", "Historical", 
                                  "Scientific", "Political", "Educational"],
                           key="genre_select")

    with col2:
        num_chapters = st.number_input(
            "Number of Chapters",
            min_value=1,
            max_value=50,
            value=templates[template]["default_chapters"],
            key="num_chapters_input"
        )
        
        words_per_chapter = st.number_input(
            "Words per Chapter",
            min_value=100,
            max_value=5000,
            value=1000,
            key="words_per_chapter_input"
        )

    # Generate intelligent chapter titles based on user prompt
    if user_prompt and st.button("Generate Chapter Titles"):
        with st.spinner("Generating chapter titles based on your prompt..."):
            chapter_titles = generate_intelligent_chapter_titles(
                user_prompt, 
                num_chapters, 
                template, 
                genre
            )
            st.session_state.generated_chapter_titles = chapter_titles
    else:
        chapter_titles = st.session_state.get('generated_chapter_titles', 
                        [f"Chapter {i+1}" for i in range(num_chapters)])

    # Display and allow editing of generated chapter titles
    st.subheader("Chapter Titles")
    edited_chapter_titles = []
    cols = st.columns(3)
    for i in range(num_chapters):
        with cols[i % 3]:
            title = st.text_input(
                f"Chapter {i+1}",
                value=chapter_titles[i] if i < len(chapter_titles) else f"Chapter {i+1}",
                key=f"chapter_title_{i}"
            )
            edited_chapter_titles.append(title)

    if st.button("Save Settings", key="save_settings_button"):
        st.session_state.book_settings = {
            "title": title,
            "template": template,
            "genre": genre,
            "num_chapters": num_chapters,
            "words_per_chapter": words_per_chapter,
            "chapter_titles": edited_chapter_titles,
            "total_words": num_chapters * words_per_chapter,
            "user_prompt": user_prompt
        }
        st.success("✅ Settings saved successfully! You can now proceed to Generate Content.")

def generate_intelligent_chapter_titles(prompt, num_chapters, template, genre):
    """Generate intelligent chapter titles based on user input"""
    # This is a placeholder for the actual AI generation
    # You would typically use an LLM API here
    system_prompt = f"""
    Based on the following user's book idea, generate {num_chapters} chapter titles.
    The book is in the {genre} genre and follows the {template} style.
    
    User's book idea:
    {prompt}
    
    Generate unique, engaging chapter titles that follow a logical progression
    and reflect the themes and elements mentioned in the user's prompt.
    """
    
    # For now, returning placeholder titles
    # Replace this with actual AI generation
    return [f"Generated Chapter {i+1}" for i in range(num_chapters)]

def display_generate_content():
    """Display the content generation page"""
    if not st.session_state.book_settings:
        st.warning("⚠️ Please configure book settings first!")
        return

    st.header("Generate Content")
    
    # Display current settings
    st.subheader("Current Settings")
    settings = st.session_state.book_settings
    st.write(f"📚 Title: {settings['title']}")
    st.write(f"📝 Style: {settings['template']}")
    st.write(f"📖 Chapters: {settings['num_chapters']}")
    st.write(f"📊 Total Words: {settings['total_words']}")
    
    if st.button("Generate Book", key="generate_book_button"):
        with st.spinner("🔄 Generating your book... This may take a few minutes."):
            progress_bar = st.progress(0)
            chapters = []
            
            for i, title in enumerate(st.session_state.book_settings["chapter_titles"]):
                chapter_info = {
                    "title": title,
                    "style": st.session_state.book_settings["template"],
                    "word_count": st.session_state.book_settings["words_per_chapter"]
                }
                content = st.session_state.book_generator.generate_chapter_content(chapter_info)
                chapters.append({"title": title, "content": content})
                progress_bar.progress((i + 1) / len(st.session_state.book_settings["chapter_titles"]))
            
            st.session_state.generated_book = {
                "title": st.session_state.book_settings["title"],
                "style": st.session_state.book_settings["template"],
                "chapters": chapters
            }
            st.success("Book generated! Proceed to Preview & Download.")

def display_preview_download():
    """Display the preview and download page"""
    if not st.session_state.generated_book:
        st.warning("⚠️ Please generate book content first!")
        return
    
    st.header("Preview & Download")
    
    # Chapter preview
    st.subheader("Chapter Preview")
    book = st.session_state.generated_book
    chapter_select = st.selectbox(
        "Select chapter to preview",
        options=[chapter["title"] for chapter in book["chapters"]],
        key="chapter_preview_select"
    )
    
    selected_chapter = next(
        chapter for chapter in book["chapters"]
        if chapter["title"] == chapter_select
    )
    
    with st.expander("Chapter Content", expanded=True):
        st.markdown(selected_chapter["content"])
    
    # Download section
    st.subheader("Download Options")
    if st.button("Generate PDF", key="generate_pdf_button"):
        with st.spinner("🔄 Creating PDF..."):
            # Here you would call your PDF generation code
            # For now, we'll just show a success message
            st.success("✅ PDF generated! Click below to download.")
            st.download_button(
                label="⬇️ Download Book",
                data=b"Sample PDF content",  # Replace with actual PDF data
                file_name=f"{book['title']}.pdf",
                mime="application/pdf"
            )

if __name__ == "__main__":
    main()