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

import pandas as pd
import streamlit as st
from PIL import Image
from reportlab.pdfgen import canvas
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 • Full-Page & Custom Layout Generator")
st.markdown(
    "Upload images, filter by orientation, reorder visually, and generate a captioned PDF where each page matches its image's dimensions."
)

# --- 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"**Preview page size:** {page_width}×{page_height} pt")

# --- Main: Upload, Filter & Reorder -------------------
st.header("2️⃣ Upload, Filter & Reorder Images")
uploaded = st.file_uploader(
    "📂 Select PNG/JPG images", type=["png","jpg","jpeg"], accept_multiple_files=True
)
ordered_files = []
if uploaded:
    records = []
    for idx, f in enumerate(uploaded):
        im = Image.open(f)
        w, h = im.size
        ar = round(w / h, 2)
        orient = "Square" if 0.9 <= ar <= 1.1 else ("Landscape" if ar > 1.1 else "Portrait")
        records.append({
            "filename": f.name,
            "width": w,
            "height": h,
            "aspect_ratio": ar,
            "orientation": orient,
            "order": idx,
        })
    df = pd.DataFrame(records)
    dims = st.sidebar.multiselect(
        "Include orientations:", options=["Landscape","Portrait","Square"],
        default=["Landscape","Portrait","Square"]
    )
    df = df[df["orientation"].isin(dims)].reset_index(drop=True)

    st.markdown("#### Image Metadata & Order")
    st.dataframe(df.style.format({"aspect_ratio": "{:.2f}"}), use_container_width=True)
    st.markdown("*Drag rows or edit the `order` column to set PDF page sequence.*")
    try:
        edited = st.experimental_data_editor(df, num_rows="fixed", use_container_width=True)
        ordered_df = edited
    except Exception:
        edited = st.data_editor(
            df,
            column_config={
                "order": st.column_config.NumberColumn(
                    "Order", min_value=0, max_value=len(df)-1
                )
            },
            hide_index=True,
            use_container_width=True,
        )
        ordered_df = edited.sort_values("order").reset_index(drop=True)

    name2file = {f.name: f for f in uploaded}
    ordered_files = [name2file[n] for n in ordered_df["filename"] if n in name2file]

# --- Utility: Clean stems -------------------------------
def clean_stem(fn: str) -> str:
    stem = os.path.splitext(fn)[0]
    return stem.replace("-", " ").replace("_", " ")

# --- PDF Creation: Image-sized + Captions --------------
def make_image_sized_pdf(images):
    buffer = io.BytesIO()
    c = canvas.Canvas(buffer)
    for idx, f in enumerate(images, start=1):
        im = Image.open(f)
        iw, ih = im.size
        # Caption height
        cap_h = 20
        page_w, page_h = iw, ih + cap_h
        c.setPageSize((page_w, page_h))
        # Draw image
        c.drawImage(ImageReader(im), 0, cap_h, iw, ih, preserveAspectRatio=True, mask='auto')
        # Draw caption
        caption = clean_stem(f.name)
        c.setFont("Helvetica", 12)
        c.drawCentredString(page_w/2, cap_h/2, caption)
        # Page number
        c.setFont("Helvetica", 8)
        c.drawRightString(page_w - 10, 10, str(idx))
        c.showPage()
    c.save()
    buffer.seek(0)
    return buffer.getvalue()

# --- Generate & Download -------------------------------
st.header("3️⃣ Generate & Download PDF with Captions")
if st.button("🖋️ Generate Captioned PDF"):
    if not ordered_files:
        st.warning("Upload and reorder at least one image.")
    else:
        # Timestamp + weekday
        now = datetime.now()
        prefix = now.strftime("%Y-%m%d-%I%M%p") + "-" + now.strftime("%a").upper()
        # Build filename from ordered stems
        stems = [clean_stem(f.name) for f in ordered_files]
        basename = " - ".join(stems)
        fname = f"{prefix}-{basename}.pdf"
        pdf_bytes = make_image_sized_pdf(ordered_files)
        st.success(f"✅ PDF ready: **{fname}**")
        st.download_button(
            "⬇️ Download PDF", data=pdf_bytes,
            file_name=fname, mime="application/pdf"
        )
        st.markdown("#### Preview of Page 1")
        try:
            import fitz
            doc = fitz.open(stream=pdf_bytes, filetype="pdf")
            pix = doc.load_page(0).get_pixmap(matrix=fitz.Matrix(1.5,1.5))
            st.image(pix.tobytes(), use_container_width=True)
        except Exception:
            st.info("Install `pymupdf` (`fitz`) for preview.")

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