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, filter/group by shape, 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_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, Inspect & Reorder -------------------
st.header("2️⃣ Upload, Inspect & Reorder Images")
uploaded = st.file_uploader(
    "📂 Select PNG/JPG images", type=["png", "jpg", "jpeg"], accept_multiple_files=True
)

if uploaded:
    # Collect metadata
    records = []
    for idx, f in enumerate(uploaded):
        im = Image.open(f)
        w, h = im.size
        ar = round(w / h, 2)
        if ar > 1.1:
            orient = "Landscape"
        elif ar < 0.9:
            orient = "Portrait"
        else:
            orient = "Square"
        records.append({
            "filename": f.name,
            "width": w,
            "height": h,
            "aspect_ratio": ar,
            "orientation": orient,
            "order": idx,
        })
    df = pd.DataFrame(records)

    # Filter by orientation
dims = st.sidebar.multiselect(
    "Include orientations:",
    options=["Landscape", "Portrait", "Square"],
    default=["Landscape", "Portrait", "Square"]
)
df = df[df["orientation"].isin(dims)]

    # Show table and allow numeric reordering
st.markdown("#### Image Metadata")
if uploaded:
    st.dataframe(df.style.format({"aspect_ratio": "{:.2f}"}), use_container_width=True)

st.markdown("#### Reorder Panels")
if uploaded:
    edited = st.data_editor(
        df,
        column_config={
            "order": st.column_config.NumberColumn(
                "Order", min_value=0, max_value=len(df) - 1
            )
        },
        hide_dataframe_index=True,
        use_container_width=True,
    )
    # sort by user-defined order
    ordered = edited.sort_values("order").reset_index(drop=True)
    name2file = {f.name: f for f in uploaded}
    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):
    buf = io.BytesIO()
    c = canvas.Canvas(buf, pagesize=(w_pt, h_pt))
    N = len(images)
    cols = int(math.ceil(math.sqrt(N)))
    rows = int(math.ceil(N / cols))
    pw = w_pt / cols
    ph = h_pt / rows
    for idx, img in enumerate(images):
        im = Image.open(img)
        iw, ih = im.size
        tar_ar = pw / ph
        img_ar = iw / ih
        if img_ar > tar_ar:
            nw = int(ih * tar_ar)
            left = (iw - nw) // 2
            im = im.crop((left, 0, left + nw, ih))
        else:
            nh = int(iw / tar_ar)
            top = (ih - nh) // 2
            im = im.crop((0, top, iw, top + nh))
        im = im.resize((int(pw), int(ph)), Image.LANCZOS)
        col = idx % cols
        row = idx // cols
        x = col * pw
        y = h_pt - (row + 1) * ph
        c.drawImage(ImageReader(im), x, y, pw, ph, preserveAspectRatio=False, mask='auto')
    c.showPage()
    c.save()
    buf.seek(0)
    return buf.getvalue()

# --- Generate & Download -------------------------------
st.header("3️⃣ Generate & Download PDF")
if st.button("🎉 Generate PDF"):
    if not ordered_files:
        st.warning("Upload and reorder at least one image.")
    else:
        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_data = make_comic_pdf(ordered_files, page_width, page_height)
        st.success(f"✅ PDF ready: **{out_name}**")
        st.download_button(
            "⬇️ Download PDF", data=pdf_data,
            file_name=out_name, mime="application/pdf"
        )
        st.markdown("#### PDF Preview")
        try:
            import fitz
            doc = fitz.open(stream=pdf_data, filetype="pdf")
            pix = doc[0].get_pixmap(matrix=fitz.Matrix(1.5, 1.5))
            st.image(pix.tobytes(), use_container_width=True)
        except Exception:
            st.info("Install `pymupdf` for live PDF preview.")

# --- Footer ------------------------------------------------
st.sidebar.markdown("---")
st.sidebar.markdown("Built by Aaron C. Wacker")