import os
import re
from datetime import datetime
from groq import Groq
from docx import Document
from docx.shared import Pt
import gradio as gr

# Initialize Groq API
client = Groq(api_key=os.environ["GROQ_API_KEY"])

# In-memory data store
book_data = {
    "title": "",
    "grade": "",
    "toc": [],
    "chapters": {},
    "preface": ""
}

# Clean formatting

def clean_formatting(text):
    text = re.sub(r"\*\*(.*?)\*\*", r"\1", text)
    text = re.sub(r"\*(.*?)\*", r"\1", text)
    text = re.sub(r"[#•●→✅🔹🔸📘📖📑📤📥⬇️🎉>]", "", text)
    text = re.sub(r"\s{2,}", " ", text)
    return text.strip()

# Generate TOC and Preface

def generate_book_intro(title, grade):
    book_data["title"] = title.strip()
    book_data["grade"] = grade.strip()

    prompt = f"""
    You are a professional educational author. 
    Write the following sections for a textbook titled: "{title}" for Grade {grade}:

    1. A formal preface aligned with APA 7th Edition style.
    2. A Table of Contents with 7 chapters. Title each chapter clearly according to a logical structure for beginners.

    Do NOT write any chapter content. Only return:
    - Preface
    - Table of Contents
    """

    response = client.chat.completions.create(
        model="llama3-70b-8192",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7
    )

    output = response.choices[0].message.content.strip()
    preface, toc = output.split("Table of Contents", 1)
    book_data["preface"] = clean_formatting(preface)
    book_data["toc"] = [clean_formatting(line) for line in toc.strip().split('\n') if line.strip()]

    return f"Preface:\n\n{book_data['preface']}\n\nTable of Contents:\n\n" + "\n".join(book_data["toc"])

# Generate chapter by number

def generate_chapter(ch_num):
    if ch_num in book_data["chapters"]:
        return book_data["chapters"][ch_num]

    ch_title = next((line for line in book_data["toc"] if line.startswith(str(ch_num))), f"Chapter {ch_num}")

    prompt = f"""
    You are an expert textbook writer. Write Chapter {ch_num} for the book titled "{book_data['title']}" for Grade {book_data['grade']}. 

    Chapter Title: {ch_title}

    Include the following sections:
    - Introduction
    - Student Learning Outcomes (5 to 7)
    - Detailed content aligned with each SLO
    - Activities or Case Studies
    - Assessment (15 MCQs with no answers)
    - Glossary

    Follow APA 7th edition tone. Do NOT use any asterisks, markdown, emojis, or non-printable characters.
    """

    response = client.chat.completions.create(
        model="llama3-70b-8192",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7
    )

    chapter = clean_formatting(response.choices[0].message.content.strip())
    book_data["chapters"][ch_num] = chapter
    return chapter

# Export full textbook to Word

def export_book():
    doc = Document()
    style = doc.styles['Normal']
    style.font.name = 'Times New Roman'
    style.font.size = Pt(12)

    def add_paragraph(text, size=12, bold=False):
        p = doc.add_paragraph()
        run = p.add_run(text.strip())
        run.font.size = Pt(size)
        run.bold = bold
        p.paragraph_format.line_spacing = 1.5

    def add_heading(text, level=1):
        size = 16 if level == 1 else 14
        add_paragraph(text.strip(), size=size, bold=True)

    # Title Page
    add_heading(book_data['title'], level=1)
    add_paragraph(f"Subject: {book_data['title']}")
    add_paragraph(f"Grade: {book_data['grade']}")
    add_paragraph("Author: AI-Generated by Groq-powered App")
    add_paragraph(f"Date: {datetime.now().strftime('%B %d, %Y')}")
    doc.add_page_break()

    # Preface
    add_heading("Preface", level=1)
    add_paragraph(book_data['preface'])
    doc.add_page_break()

    # TOC
    add_heading("Table of Contents", level=1)
    for line in book_data['toc']:
        add_paragraph(line)
    doc.add_page_break()

    # Chapters
    for num in sorted(book_data['chapters'].keys()):
        add_heading(f"{book_data['toc'][num-1]}", level=1)
        for line in book_data['chapters'][num].split('\n'):
            if any(keyword in line.lower() for keyword in ["introduction", "summary", "examples", "review", "student learning outcomes", "assessment", "glossary", "activities"]):
                add_heading(line, level=2)
            else:
                add_paragraph(line)
        doc.add_page_break()

    file_path = "/tmp/Textbook_AI_Generated.docx"
    doc.save(file_path)
    return file_path

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## 📘 AI Textbook Generator – APA Style | Developed by Najaf Ali Sharqi")

    with gr.Row():
        title = gr.Textbox(label="Subject Title")
        grade = gr.Textbox(label="Grade Level")

    generate_structure = gr.Button("🧠 Generate TOC + Preface")
    structure_output = gr.Textbox(label="Generated Preface and TOC", lines=15)

    generate_structure.click(fn=generate_book_intro, inputs=[title, grade], outputs=structure_output)

    gr.Markdown("### 📥 Generate Chapters One by One")
    chapter_output = gr.Textbox(label="Generated Chapter Content", lines=25)

    with gr.Row():
        for i in range(1, 8):
            gr.Button(f"Generate Chapter {i}").click(fn=generate_chapter, inputs=gr.Number(value=i, visible=False), outputs=chapter_output)

    gr.Markdown("### 💾 Download Complete Textbook")
    download_btn = gr.Button("Download MS Word Textbook")
    download_file = gr.File()

    download_btn.click(fn=export_book, outputs=download_file)

    gr.Markdown("Created with ❤️ using Gradio + Groq | All content APA 7th aligned")

demo.launch(share=True)