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, scan documents, 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: Scan, Upload, Filter & Reorder --------------
st.header("2️⃣ Document Scan & Image Upload")
# Document scan via camera
cam_img = st.camera_input("📸 Scan Document", key="doc_scan")
if cam_img:
    now = datetime.now()
    prefix = now.strftime("%Y-%m%d-%I%M%p") + "-" + now.strftime("%a").upper()
    scan_name = f"{prefix}-scan.png"
    with open(scan_name, "wb") as f:
        f.write(cam_img.getvalue())
    st.image(Image.open(scan_name), caption=scan_name, use_container_width=True)

# Image upload
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
        cap_h = 20
        page_w, page_h = iw, ih + cap_h
        c.setPageSize((page_w, page_h))
        c.drawImage(ImageReader(im), 0, cap_h, iw, ih, preserveAspectRatio=True, mask='auto')
        caption = clean_stem(f.name)
        c.setFont("Helvetica", 12)
        c.drawCentredString(page_w/2, cap_h/2, caption)
        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:
        now = datetime.now()
        prefix = now.strftime("%Y-%m%d-%I%M%p") + "-" + now.strftime("%a").upper()
        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")