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 • Comic-Book Layout Generator")
st.markdown(
    "Upload images, choose aspect ratio, filter by shape, reorder visually, and generate a PDF."
)

# --- 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:
    # ➊ Collect metadata
    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)

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

    # ➌ Show metadata table
    st.markdown("#### Image Metadata")
    st.dataframe(df.style.format({"aspect_ratio": "{:.2f}"}), use_container_width=True)

    # ➍ Visual reorder via Data Editor
    st.markdown("#### Reorder Panels (drag rows or adjust 'order')")
    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)

    # Map filenames to ordered files
    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_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, ph = w_pt/cols, h_pt/rows
    for idx, f in enumerate(images):
        im = Image.open(f)
        iw, ih = im.size
        tar = pw/ph; ar = iw/ih
        if ar>tar:
            nw = int(ih*tar); left=(iw-nw)//2; im=im.crop((left,0,left+nw,ih))
        else:
            nh = int(iw/tar); top=(ih-nh)//2; im=im.crop((0,top,iw,top+nh))
        im = im.resize((int(pw),int(ph)), Image.LANCZOS)
        col, row = idx%cols, idx//cols
        x, y = col*pw, 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_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_comic_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")