import io
import os
import math
import re
from collections import Counter
from datetime import datetime

import pandas as pd
import streamlit as st
from PIL import Image
from reportlab.pdfgen import canvas
from reportlab.lib.units import inch
from reportlab.lib.utils import ImageReader

# --- App Configuration ----------------------------------
st.set_page_config(
    page_title="Image → PDF Comic Layout",
    layout="wide",
    initial_sidebar_state="expanded",
)

st.title("🖼️ Image → PDF • Comic-Book Layout Generator")
st.markdown(
    "Upload images, choose a page aspect ratio, reorder panels, and generate a high‑definition PDF with smart naming."
)


# --- Sidebar: Page Settings -----------------------------
st.sidebar.header("1️⃣ Page Aspect Ratio & Size")
ratio_map = {
    "4:3 (Landscape)": (4, 3),
    "16:9 (Landscape)": (16, 9),
    "1:1 (Square)": (1, 1),
    "2:3 (Portrait)": (2, 3),
    "9:16 (Portrait)": (9, 16),
}
ratio_choice = st.sidebar.selectbox(
    "Preset Ratio", list(ratio_map.keys()) + ["Custom…"]
)
if ratio_choice != "Custom…":
    rw, rh = ratio_map[ratio_choice]
else:
    rw = st.sidebar.number_input("Custom Width Ratio", min_value=1, value=4)
    rh = st.sidebar.number_input("Custom Height Ratio", min_value=1, value=3)

# Base page width in points (1pt = 1/72 inch)
BASE_WIDTH_PT = st.sidebar.slider(
    "Base Page Width (pt)", min_value=400, max_value=1200, value=800, step=100
)
page_width = BASE_WIDTH_PT
page_height = int(BASE_WIDTH_PT * (rh / rw))

st.sidebar.markdown(f"**Page size:** {page_width}×{page_height} pt")


# --- Main: Upload & Reorder -----------------------------
st.header("2️⃣ Upload & Reorder Images")
uploaded_files = st.file_uploader(
    "📂 Select PNG/JPG images", type=["png", "jpg", "jpeg"], accept_multiple_files=True
)

# Build ordering table
if uploaded_files:
    df = pd.DataFrame({"filename": [f.name for f in uploaded_files]})
    st.markdown("Drag to reorder panels below:")
    ordered = st.experimental_data_editor(
        df, num_rows="fixed", use_container_width=True
    )
    # Map back to actual file objects in new order
    name2file = {f.name: f for f in uploaded_files}
    ordered_files = [name2file[n] for n in ordered["filename"] if n in name2file]
else:
    ordered_files = []


# --- PDF Creation Logic ----------------------------------

def top_n_words(filenames, n=5):
    words = []
    for fn in filenames:
        stem = os.path.splitext(fn)[0]
        words += re.findall(r"\w+", stem.lower())
    return [w for w, _ in Counter(words).most_common(n)]


def make_comic_pdf(images, w_pt, h_pt):
    buffer = io.BytesIO()
    c = canvas.Canvas(buffer, pagesize=(w_pt, h_pt))

    N = len(images)
    cols = int(math.ceil(math.sqrt(N)))
    rows = int(math.ceil(N / cols))
    panel_w = w_pt  / cols
    panel_h = h_pt / rows

    for idx, img_file in enumerate(images):
        im = Image.open(img_file)
        iw, ih = im.size
        target_ar = panel_w / panel_h
        img_ar = iw / ih

        # Center-crop to panel aspect
        if img_ar > target_ar:
            new_w = int(ih * target_ar)
            left = (iw - new_w) // 2
            im = im.crop((left, 0, left + new_w, ih))
        else:
            new_h = int(iw / target_ar)
            top = (ih - new_h) // 2
            im = im.crop((0, top, iw, top + new_h))

        im = im.resize((int(panel_w), int(panel_h)), Image.LANCZOS)

        col = idx % cols
        row = idx // cols
        x = col * panel_w
        y = h_pt - (row + 1) * panel_h

        c.drawImage(
            ImageReader(im), x, y, panel_w, panel_h,
            preserveAspectRatio=False, mask='auto'
        )

    c.showPage()
    c.save()
    buffer.seek(0)
    return buffer.getvalue()


# --- Generate & Download -------------------------------
st.header("3️⃣ Generate & Download PDF")
if st.button("🎉 Generate PDF"):
    if not ordered_files:
        st.warning("Please upload and order at least one image.")
    else:
        # Build filename: YYYY-MMdd-top5words.pdf
        date_str = datetime.now().strftime("%Y-%m%d")
        words = top_n_words([f.name for f in ordered_files], n=5)
        slug  = "-".join(words)
        out_name = f"{date_str}-{slug}.pdf"

        pdf_bytes = make_comic_pdf(ordered_files, page_width, page_height)

        st.success(f"✅ PDF ready: **{out_name}**")
        st.download_button(
            "⬇️ Download PDF", data=pdf_bytes,
            file_name=out_name, mime="application/pdf"
        )

        # Preview first page (requires pymupdf)
        st.markdown("#### PDF Preview")
        try:
            import fitz  # pymupdf
            doc = fitz.open(stream=pdf_bytes, filetype="pdf")
            pix = doc[0].get_pixmap(matrix=fitz.Matrix(1.5, 1.5))
            st.image(pix.tobytes(), use_column_width=True)
        except Exception:
            st.info("Install `pymupdf` for live PDF preview.")


# --- Footer ------------------------------------------------
st.sidebar.markdown("---")
st.sidebar.markdown(
    "Built by Aaron C. Wacker • Senior AI Engineer"
)