import gradio as gr
import os
import tempfile

# For Word processing
try:
    from docx import Document
    from docx.shared import Inches
    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False

# For PDF processing
try:
    import PyPDF2
    from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
    from reportlab.lib.pagesizes import inch
    from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
    from reportlab.lib.enums import TA_CENTER
    REPORTLAB_AVAILABLE = True
except ImportError:
    REPORTLAB_AVAILABLE = False

# ---------------- KDP Spec Config ---------------------
TRIM_SIZES = [
    {"size": "5\" x 8\"", "width": 5.0, "height": 8.0, "min": 24, "max": 828},
    {"size": "5.06\" x 7.81\"", "width": 5.06, "height": 7.81, "min": 24, "max": 828},
    {"size": "5.25\" x 8\"", "width": 5.25, "height": 8.0, "min": 24, "max": 828},
    {"size": "5.5\" x 8.5\"", "width": 5.5, "height": 8.5, "min": 24, "max": 828},
    {"size": "6\" x 9\"", "width": 6.0, "height": 9.0, "min": 24, "max": 828},

    # Large (cost more)
    {"size": "6.14\" x 9.21\"", "width": 6.14, "height": 9.21, "min": 24, "max": 550},
    {"size": "6.69\" x 9.61\"", "width": 6.69, "height": 9.61, "min": 24, "max": 550},
    {"size": "7\" x 10\"", "width": 7.0, "height": 10.0, "min": 24, "max": 550},
    {"size": "7.44\" x 9.69\"", "width": 7.44, "height": 9.69, "min": 24, "max": 550},
    {"size": "7.5\" x 9.25\"", "width": 7.5, "height": 9.25, "min": 24, "max": 550},
    {"size": "8\" x 10\"", "width": 8.0, "height": 10.0, "min": 24, "max": 550},
    {"size": "8.25\" x 6\"", "width": 8.25, "height": 6.0, "min": 24, "max": 550},
    {"size": "8.25\" x 8.25\"", "width": 8.25, "height": 8.25, "min": 24, "max": 550},
    {"size": "8.5\" x 8.5\"", "width": 8.5, "height": 8.5, "min": 24, "max": 550},
    {"size": "8.5\" x 11\"", "width": 8.5, "height": 11.0, "min": 24, "max": 550},
    {"size": "8.27\" x 11.69\"", "width": 8.27, "height": 11.69, "min": 24, "max": 550},
]

def get_trim(trim_label):
    for trim in TRIM_SIZES:
        if trim_label.startswith(trim['size']):
            return trim
    return None

def get_margin(page_count, has_bleed):
    """Return inside, outside (left), top, bottom margins in inches"""
    if page_count <= 150:
        inside = 0.375
    elif page_count <= 300:
        inside = 0.5
    elif page_count <= 500:
        inside = 0.625
    elif page_count <= 700:
        inside = 0.75
    else:
        inside = 0.875
    outside = 0.375 if has_bleed else 0.25
    top = bottom = 0.375 if has_bleed else 0.25
    return inside, outside, top, bottom

def get_bleed_dimensions(width, height, has_bleed):
    if has_bleed:
        return width + 0.25, height + 0.25
    return width, height

def spine_width(page_count, paper_type, ink):
    if ink == "color":
        return page_count * 0.002347
    elif paper_type == "cream":
        return page_count * 0.0025
    else:
        return page_count * 0.002252

def trim_choices():
    return [
        f'{trim["size"]} ({"Regular" if trim["max"]==828 else "Large"})'
        for trim in TRIM_SIZES
    ]

# -------------------- Core Formatting Logic -------------------------

def format_word(file_path, trim_size, has_bleed):
    try:
        if not DOCX_AVAILABLE:
            return "python-docx library is required (pip install python-docx)", None

        trim = get_trim(trim_size)
        if not trim:
            return f"Invalid trim size: {trim_size}", None

        doc = Document(file_path)

        text_len = sum(len(p.text) for p in doc.paragraphs)
        page_count = max(24, int(text_len / 2000) or 1)

        width, height = get_bleed_dimensions(trim['width'], trim['height'], has_bleed)
        inside, outside, top, bottom = get_margin(page_count, has_bleed)

        for section in doc.sections:
            section.page_width = Inches(width)
            section.page_height = Inches(height)
            section.left_margin = Inches(inside)
            section.right_margin = Inches(outside)
            section.top_margin = Inches(top)
            section.bottom_margin = Inches(bottom)

        out_path = file_path.replace(".docx", "_kdp.docx")
        doc.save(out_path)
        return (
            f"SUCCESS!\nTrim size: {trim_size}\nPage count (approx): {page_count}\nMargins (in): Inside={inside}, Out={outside}, Top/Bottom={top}\nBleed: {'Yes' if has_bleed else 'No'}",
            out_path
        )
    except Exception as e:
        return f"Error processing Word document: {str(e)}", None

def format_pdf(file_path, trim_size, has_bleed):
    try:
        if not REPORTLAB_AVAILABLE:
            return "reportlab and PyPDF2 are required", None

        trim = get_trim(trim_size)
        if not trim:
            return f"Invalid trim size: {trim_size}", None

        with open(file_path, "rb") as f:
            pdf = PyPDF2.PdfReader(f)
            page_count = len(pdf.pages)

        width, height = get_bleed_dimensions(trim['width'], trim['height'], has_bleed)
        inside, outside, top, bottom = get_margin(page_count, has_bleed)
        out_path = file_path.replace('.pdf', '_kdp.pdf')

        # Create formatted PDF with specifications
        doc = SimpleDocTemplate(
            out_path,
            pagesize=(width * inch, height * inch),
            leftMargin=inside * inch,
            rightMargin=outside * inch,
            topMargin=top * inch,
            bottomMargin=bottom * inch
        )
        styles = getSampleStyleSheet()
        story = [
            Paragraph("Amazon KDP Formatting Summary", ParagraphStyle('Heading1', fontSize=20, alignment=TA_CENTER)),
            Spacer(1, 12),
            Paragraph(
                f"<b>Trim size:</b> {trim_size}<br>"
                f"<b>Pages:</b> {page_count}<br>"
                f"<b>Page dimensions (in):</b> {width} x {height}<br>"
                f"<b>Margins:</b> Inside={inside}, Out={outside}, Top/Bottom={top}<br>"
                f"<b>Bleed:</b> {'Yes' if has_bleed else 'No'}<br>",
                styles['BodyText']
            )
        ]
        doc.build(story)
        return f"SUCCESS! PDF processed and formatted for KDP!\nPages: {page_count}\nMargins applied for KDP compliance.", out_path
    except Exception as e:
        return f"Error processing PDF: {str(e)}", None

def format_document(file, trim_size, has_bleed, paper_type, ink_type):
    try:
        if file is None:
            return "Please upload a .docx or .pdf file.", None

        # Handle both file objects and file paths (Gradio can pass either)
        if hasattr(file, 'name') and hasattr(file, 'read'):
            # File object from Gradio
            file_path = file.name
            ext = os.path.splitext(file_path)[1].lower()

            # Create temporary file
            with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
                tmp.write(file.read())
                tmp_file_path = tmp.name
        elif isinstance(file, str):
            # File path string
            tmp_file_path = file
            ext = os.path.splitext(file)[1].lower()
        else:
            return "Invalid file input", None

        if ext == '.docx':
            result, out_file = format_word(tmp_file_path, trim_size, has_bleed)
        elif ext == '.pdf':
            result, out_file = format_pdf(tmp_file_path, trim_size, has_bleed)
        else:
            return "Only .docx and .pdf files are supported", None

        # Clean up temporary file if we created one
        if hasattr(file, 'read') and os.path.exists(tmp_file_path):
            os.unlink(tmp_file_path)

        return result, out_file
    except Exception as e:
        return f"Error in format_document: {str(e)}", None

def calculate_specs(trim_size, page_count, paper_type, ink_type):
    try:
        trim = get_trim(trim_size)
        if not trim:
            return "Select a valid trim size."
        if not (trim["min"] <= page_count <= trim["max"]):
            return f"Invalid page count for {trim['size']}: must be {trim['min']}-{trim['max']}"
        inside, outside, top, bottom = get_margin(page_count, False)
        inside_b, outside_b, top_b, bottom_b = get_margin(page_count, True)
        width, height = trim['width'], trim['height']
        width_b, height_b = get_bleed_dimensions(width, height, True)
        spine = spine_width(page_count, paper_type, ink_type)
        size_type = "Regular" if trim["max"] == 828 else "Large"
        return (
            f"<b>Trim size:</b> {trim['size']} ({size_type})<br>"
            f"<b>Pages:</b> {page_count}<br>"
            f"<b>Spine width:</b> {spine:.4f}\"<br>"
            f"<b>Page dimensions:</b> {width}\" x {height}\"<br>"
            f"<b>Page dimensions (w/bleed):</b> {width_b}\" x {height_b}\"<br>"
            "<b>Margins (no bleed):</b> " +
            f"Inside {inside}\", Out {outside}\", Top/Bottom {top}\"<br>"
            "<b>Margins (with bleed):</b> " +
            f"Inside {inside_b}\", Out {outside_b}\", Top/Bottom {top_b}\"<br>"
            f"<b>Paper:</b> {paper_type}, <b>Ink:</b> {ink_type}"
        )
    except Exception as e:
        return f"Error calculating specs: {str(e)}"

# -------------------- Gradio Interface -----------------------
with gr.Blocks(title="Amazon KDP Paperback Formatter") as demo:
    gr.Markdown("""
    # 📚 Amazon KDP Paperbook Formatter

    Upload a .docx or .pdf manuscript and choose your KDP trim size.
    Set bleed (for graphics to the page edge), paper type and ink.
    Download your ready-to-upload file.
    """)
    with gr.Tab("📝 Format Document"):
        with gr.Row():
            with gr.Column():
                file_input = gr.File(label="Upload .docx or .pdf", file_types=['.docx', '.pdf'])
                trim_size_dropdown = gr.Dropdown(
                    choices=trim_choices(), label="Trim Size",
                    value="6\" x 9\" (Regular)"
                )
                has_bleed_checkbox = gr.Checkbox(label="Include Bleed", value=False)
                paper_type_radio = gr.Radio(choices=["white", "cream"], label="Paper", value="white")
                ink_type_radio = gr.Radio(choices=["black", "color"], label="Ink", value="black")
                format_button = gr.Button("Format for KDP", variant="primary")
            with gr.Column():
                result_output = gr.Textbox(label="Result", lines=8)
                download_output = gr.File(label="Download", visible=False)

    with gr.Tab("📊 Calculate Specs"):
        calc_trim = gr.Dropdown(trim_choices(), label="Trim Size", value="6\" x 9\" (Regular)")
        calc_page = gr.Slider(24, 828, value=100, step=1, label="Page Count")
        calc_paper = gr.Radio(["white", "cream"], label="Paper", value="white")
        calc_ink = gr.Radio(["black", "color"], label="Ink", value="black")
        calc_output = gr.HTML(label="Specifications")
        calc_btn = gr.Button("Calculate")

    with gr.Tab("ℹ️ KDP Help"):
        gr.Markdown("""
        ## Amazon KDP Paperback Guidelines

        **Popular Trim Sizes:**
        - **6" x 9"** ⭐ Most popular for novels/non-fiction
        - **5.5" x 8.5"** Good for smaller books
        - **8.5" x 11"** Workbooks, journals, reference

        **Margins by Page Count:**
        - 24-150 pages: 0.375" inside margin
        - 151-300 pages: 0.5" inside margin  
        - 301-500 pages: 0.625" inside margin
        - 501-700 pages: 0.75" inside margin
        - 701-828 pages: 0.875" inside margin

        **Outside/Top/Bottom Margins:**
        - Without bleed: 0.25"
        - With bleed: 0.375"

        **Bleed:**
        - Add 0.125" to all edges when graphics extend to page edge
        - Example: 6" x 9" becomes 6.125" x 9.25" with bleed

        **Spine Width:**
        - B&W (White): pages × 0.002252"
        - B&W (Cream): pages × 0.0025" 
        - Color: pages × 0.002347"

        **Page Limits:**
        - Regular sizes: 24-828 pages
        - Large sizes: 24-550 pages
        """)

    # Event handlers
    format_button.click(
        fn=format_document,
        inputs=[file_input, trim_size_dropdown, has_bleed_checkbox, paper_type_radio, ink_type_radio],
        outputs=[result_output, download_output]
    ).then(
        fn=lambda x: gr.update(visible=x is not None),
        inputs=[download_output], 
        outputs=[download_output]
    )

    calc_btn.click(
        fn=calculate_specs,
        inputs=[calc_trim, calc_page, calc_paper, calc_ink],
        outputs=[calc_output]
    )

    # Auto-calculate on changes
    for component in [calc_trim, calc_page, calc_paper, calc_ink]:
        component.change(
            fn=calculate_specs,
            inputs=[calc_trim, calc_page, calc_paper, calc_ink],
            outputs=[calc_output]
        )

if __name__ == "__main__":
    demo.launch()