# app.py
import tempfile
from pathlib import Path
from typing import Tuple

import gradio as gr
from docx import Document
from docx.shared import Inches
from pypdf import PdfReader
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.pagesizes import inch
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont

# Register a standard font (optional)
# pdfmetrics.registerFont(TTFont('DejaVu', 'DejaVuSans.ttf'))

# --- KDP Constants ---
TRIM_SIZES_INCHES = {
    "5\" × 8\"": (5.0, 8.0),
    "5.5\" × 8.5\"": (5.5, 8.5),
    "6\" × 9\" (Most Popular)": (6.0, 9.0),
    "7\" × 10\"": (7.0, 10.0),
    "8.5\" × 11\"": (8.5, 11.0),
}

GUTTER_MARGINS_INCHES = [
    (150, 0.375),
    (300, 0.5),
    (500, 0.625),
    (700, 0.75),
    (828, 0.875),
]

def extract_text_from_docx(docx_path: str) -> list:
    """Extract paragraphs from DOCX."""
    doc = Document(docx_path)
    return [para.text for para in doc.paragraphs if para.text.strip()]

def extract_text_from_pdf(pdf_path: str) -> list:
    """Extract text from PDF (simple)."""
    reader = PdfReader(pdf_path)
    text = []
    for page in reader.pages:
        page_text = page.extract_text()
        if page_text:
            text.extend([line for line in page_text.split("\n") if line.strip()])
    return text

def apply_kdp_to_docx(docx_path: str, trim_w: float, trim_h: float, page_count: int, has_bleed: bool) -> str:
    doc = Document(docx_path)
    section = doc.sections[0]
    
    section.page_width = Inches(trim_w)
    section.page_height = Inches(trim_h)
    
    outside = 0.375 if has_bleed else 0.25
    gutter = 0.375
    for max_p, g in GUTTER_MARGINS_INCHES:
        if page_count <= max_p:
            gutter = g
            break

    section.top_margin = Inches(outside)
    section.bottom_margin = Inches(outside)
    section.left_margin = Inches(gutter)
    section.right_margin = Inches(outside)

    output_path = Path(tempfile.gettempdir()) / f"KDP_formatted_{Path(docx_path).stem}.docx"
    doc.save(output_path)
    return str(output_path)

def create_pdf_from_text(text_lines: list, output_path: str, trim_w: float, trim_h: float, has_bleed: bool):
    """Create a simple PDF from text lines."""
    page_size = (trim_w * inch, trim_h * inch)
    outside_margin = 0.375 if has_bleed else 0.25
    
    doc = SimpleDocTemplate(
        output_path,
        pagesize=page_size,
        leftMargin=gutter_to_inch(len(text_lines)) * inch,  # Approximate
        rightMargin=outside_margin * inch,
        topMargin=outside_margin * inch,
        bottomMargin=outside_margin * inch
    )
    
    styles = getSampleStyleSheet()
    story = []
    for line in text_lines:
        story.append(Paragraph(line, styles['Normal']))
        story.append(Spacer(1, 6))
    
    doc.build(story)

def gutter_to_inch(page_count: int) -> float:
    gutter = 0.375
    for max_p, g in GUTTER_MARGINS_INCHES:
        if page_count <= max_p:
            gutter = g
            break
    return gutter

def process_file(
    input_file: str,
    trim_size: str,
    has_bleed: bool,
    page_count: int,
) -> Tuple[str, str, str]:
    file_path = Path(input_file)
    trim_w, trim_h = TRIM_SIZES_INCHES[trim_size]
    
    # Step 1: Extract text and get a working DOCX
    if file_path.suffix.lower() == ".pdf":
        text_lines = extract_text_from_pdf(str(file_path))
        # Create a new DOCX from text
        temp_doc = Document()
        for line in text_lines:
            temp_doc.add_paragraph(line)
        docx_path = Path(tempfile.gettempdir()) / "temp_from_pdf.docx"
        temp_doc.save(docx_path)
    else:
        text_lines = extract_text_from_docx(str(file_path))
        docx_path = file_path

    # Step 2: Format DOCX
    formatted_docx = apply_kdp_to_docx(str(docx_path), trim_w, trim_h, page_count, has_bleed)

    # Step 3: Create PDF from text
    formatted_pdf = Path(tempfile.gettempdir()) / f"KDP_formatted_{file_path.stem}.pdf"
    create_pdf_from_text(text_lines, str(formatted_pdf), trim_w, trim_h, has_bleed)

    info = f"Page Count Used: {page_count}\n⚠️ PDF is text-only (no images/formatting)."
    return formatted_docx, str(formatted_pdf), info

# --- UI ---
with gr.Blocks() as demo:
    gr.Markdown("# 📚 KDP Formatter (HF Spaces Compatible)")
    gr.Markdown("Upload `.docx` or `.pdf`. **PDF output is text-only** due to platform limits.")
    
    with gr.Row():
        with gr.Column():
            input_file = gr.File(file_types=[".docx", ".pdf"])
            trim_size = gr.Dropdown(choices=list(TRIM_SIZES_INCHES.keys()), value="6\" × 9\" (Most Popular)")
            page_count = gr.Number(value=300, minimum=24, maximum=828, label="Page Count")
            has_bleed = gr.Checkbox(label="Include Bleed")
            submit = gr.Button("Format")
        with gr.Column():
            out_docx = gr.File(label="Formatted DOCX")
            out_pdf = gr.File(label="Formatted PDF (Text Only)")
            info = gr.Textbox(label="Info")

    submit.click(
        process_file,
        [input_file, trim_size, has_bleed, page_count],
        [out_docx, out_pdf, info]
    )

    gr.Markdown("💡 For full PDF formatting, run this app locally with LibreOffice installed.")

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