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.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 Layout Generator")
st.markdown(
    "Upload images, filter by orientation, reorder visually, and generate a PDF with each image filling its own page."
)

# --- 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, 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")
    st.dataframe(df.style.format({"aspect_ratio": "{:.2f}"}), use_container_width=True)

    st.markdown("#### Reorder Panels")
    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]

# --- PDF Creation Logic ----------------------------------
def top_n_words(names, n=5):
    words = []
    for fn in names:
        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_fullpage_pdf(images, w_pt, h_pt):
    buf = io.BytesIO()
    c = canvas.Canvas(buf, pagesize=(w_pt, h_ht := h_pt))
    for f in images:
        im = Image.open(f)
        iw, ih = im.size
        # calculate scale to fit page
        scale = min(w_pt/iw, h_pt/ih)
        new_w = int(iw * scale)
        new_h = int(ih * scale)
        x = (w_pt - new_w)/2
        y = (h_pt - new_h)/2
        im = im.resize((new_w, new_h), Image.LANCZOS)
        c.drawImage(ImageReader(im), x, y, new_w, new_h, 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_s = datetime.now().strftime("%Y-%m%d")
        words = top_n_words([f.name for f in ordered_files])
        slug = "-".join(words)
        fname = f"{date_s}-{slug}.pdf"
        pdf = make_fullpage_pdf(ordered_files, page_width, page_height)
        st.success(f"✅ PDF ready: **{fname}**")
        st.download_button("⬇️ Download PDF", data=pdf, file_name=fname, mime="application/pdf")
        st.markdown("#### Preview")
        try:
            import fitz
            doc = fitz.open(stream=pdf, filetype="pdf")
            pix = doc[0].get_pixmap(matrix=fitz.Matrix(1.5,1.5))
            st.image(pix.tobytes(), use_container_width=True)
        except:
            st.info("Install `pymupdf` for preview.")

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