import openai
import time
from docx import Document
from docx.shared import Inches
import os
from datetime import datetime

# Configuration
OPENAI_API_KEY = "your-openai-api-key-here"  # Replace with your actual API key
MODEL = "gpt-4"  # or "gpt-3.5-turbo" for cheaper option
DELAY_BETWEEN_REQUESTS = 2  # seconds to avoid rate limiting

class StoryGenerator:
    def __init__(self, api_key):
        self.client = openai.OpenAI(api_key=api_key)
        self.story_outline = ""
        self.chapter_summaries = []
        self.chapters = []

    def generate_story_outline(self, story_type="hero's journey", genre="fantasy", additional_notes=""):
        """Generate initial story outline"""
        prompt = f"""
        Create a detailed story outline using the {story_type} storytelling arc in the {genre} genre.

        Include:
        - Basic character development (protagonist, antagonist, supporting characters)
        - World building elements
        - Major plot points following the {story_type} structure
        - Themes and character arcs
        - Setting and atmosphere

        Additional requirements: {additional_notes}

        This should be a complete outline for a novella-length story that can be divided into 12 chapters.
        """

        print("Generating story outline...")
        response = self.client.chat.completions.create(
            model=MODEL,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=1500,
            temperature=0.7
        )

        self.story_outline = response.choices[0].message.content
        print("✓ Story outline generated")
        return self.story_outline

    def organize_into_chapters(self):
        """Organize the story outline into 12 chapters"""
        prompt = f"""
        Based on this story outline:

        {self.story_outline}

        Organize this story into exactly 12 chapters. For each chapter, provide:
        - Chapter number and title
        - 2-3 sentence summary of what happens in that chapter
        - Key characters involved
        - Important plot points or developments

        Format each chapter summary clearly and number them 1-12.
        """

        print("Organizing story into 12 chapters...")
        time.sleep(DELAY_BETWEEN_REQUESTS)

        response = self.client.chat.completions.create(
            model=MODEL,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=1500,
            temperature=0.7
        )

        chapter_organization = response.choices[0].message.content
        print("✓ Story organized into chapters")
        return chapter_organization

    def extract_chapter_summaries(self, chapter_organization):
        """Extract individual chapter summaries for later use"""
        # This is a simple extraction - you might want to make this more robust
        lines = chapter_organization.split('\n')
        current_chapter = ""

        for line in lines:
            if line.strip().startswith(('Chapter', 'CHAPTER')) and any(str(i) in line for i in range(1, 13)):
                if current_chapter:
                    self.chapter_summaries.append(current_chapter.strip())
                current_chapter = line + '\n'
            elif current_chapter and line.strip():
                current_chapter += line + '\n'

        if current_chapter:
            self.chapter_summaries.append(current_chapter.strip())

        print(f"✓ Extracted {len(self.chapter_summaries)} chapter summaries")

    def write_chapter(self, chapter_number, chapter_summary):
        """Write a single chapter based on its summary"""
        prompt = f"""
        Based on this story outline:
        {self.story_outline}

        Write Chapter {chapter_number} with the following summary:
        {chapter_summary}

        Write this as a complete chapter of approximately 800-1200 words. Include:
        - Engaging prose and dialogue
        - Character development
        - Atmospheric descriptions
        - Smooth narrative flow
        - Proper chapter pacing

        Write in a engaging, literary style appropriate for the genre.
        """

        print(f"Writing Chapter {chapter_number}...")
        time.sleep(DELAY_BETWEEN_REQUESTS)

        response = self.client.chat.completions.create(
            model=MODEL,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=1500,
            temperature=0.8
        )

        chapter_content = response.choices[0].message.content
        self.chapters.append({
            'number': chapter_number,
            'content': chapter_content,
            'summary': chapter_summary
        })

        print(f"✓ Chapter {chapter_number} completed ({len(chapter_content)} characters)")
        return chapter_content

    def write_all_chapters(self):
        """Write all 12 chapters"""
        for i, summary in enumerate(self.chapter_summaries, 1):
            try:
                self.write_chapter(i, summary)
            except Exception as e:
                print(f"Error writing chapter {i}: {e}")
                continue

    def save_to_word_document(self, filename=None):
        """Save the complete story to a Word document"""
        if not filename:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"ai_generated_story_{timestamp}.docx"

        doc = Document()

        # Add title page
        title_para = doc.add_heading('AI Generated Story', 0)
        title_para.alignment = 1  # Center alignment

        doc.add_paragraph(f"Generated on: {datetime.now().strftime('%B %d, %Y')}")
        doc.add_paragraph(f"Total Chapters: {len(self.chapters)}")
        doc.add_page_break()

        # Add story outline
        doc.add_heading('Story Outline', level=1)
        doc.add_paragraph(self.story_outline)
        doc.add_page_break()

        # Add each chapter
        for chapter in self.chapters:
            doc.add_heading(f"Chapter {chapter['number']}", level=1)
            doc.add_paragraph(chapter['content'])
            doc.add_page_break()

        doc.save(filename)
        print(f"✓ Story saved to {filename}")
        return filename

def main():
    # Check if API key is set
    if OPENAI_API_KEY == "your-openai-api-key-here":
        print("❌ Please set your OpenAI API key in the OPENAI_API_KEY variable")
        return

    # Initialize the story generator
    generator = StoryGenerator(OPENAI_API_KEY)

    # Get user preferences
    print("=== AI Story Generator ===\n")

    story_type = input("Enter story arc type (default: hero's journey): ").strip() or "hero's journey"
    genre = input("Enter genre (default: fantasy): ").strip() or "fantasy"
    additional_notes = input("Additional requirements (optional): ").strip()

    try:
        # Step 1: Generate story outline
        outline = generator.generate_story_outline(story_type, genre, additional_notes)
        print("\n" + "="*50)
        print("STORY OUTLINE:")
        print("="*50)
        print(outline)

        # Ask if user wants to continue or regenerate
        while True:
            choice = input("\nAre you satisfied with this outline? (y/n/regenerate): ").lower()
            if choice in ['y', 'yes']:
                break
            elif choice in ['regenerate', 'r']:
                outline = generator.generate_story_outline(story_type, genre, additional_notes)
                print("\n" + "="*50)
                print("REGENERATED STORY OUTLINE:")
                print("="*50)
                print(outline)
            elif choice in ['n', 'no']:
                print("Exiting...")
                return

        # Step 2: Organize into chapters
        chapter_org = generator.organize_into_chapters()
        generator.extract_chapter_summaries(chapter_org)

        print("\n" + "="*50)
        print("CHAPTER ORGANIZATION:")
        print("="*50)
        print(chapter_org)

        # Ask if user wants to proceed with writing
        proceed = input(f"\nProceed to write all {len(generator.chapter_summaries)} chapters? This may take several minutes. (y/n): ")
        if proceed.lower() not in ['y', 'yes']:
            print("Stopping before chapter generation.")
            return

        # Steps 3-14: Write all chapters
        print("\n" + "="*50)
        print("GENERATING CHAPTERS...")
        print("="*50)

        generator.write_all_chapters()

        # Save to Word document
        filename = generator.save_to_word_document()

        print("\n" + "="*50)
        print("GENERATION COMPLETE!")
        print("="*50)
        print(f"Total chapters written: {len(generator.chapters)}")
        print(f"Story saved to: {filename}")

        # Calculate approximate word count
        total_chars = sum(len(chapter['content']) for chapter in generator.chapters)
        approx_words = total_chars // 5  # Rough estimate
        print(f"Approximate word count: {approx_words:,} words")

    except Exception as e:
        print(f"❌ An error occurred: {e}")
        print("Make sure your OpenAI API key is valid and you have sufficient credits.")

if __name__ == "__main__":
    main()