app.py

import streamlit as st
from docx import Document
from docx.shared import Inches, Pt
from docx.enum.table import WD_TABLE_ALIGNMENT, WD_ALIGN_VERTICAL, WD_ROW_HEIGHT_RULE
from docx.enum.text import WD_ALIGN_PARAGRAPH
from PIL import Image
import pillow_heif
import os
import io
import subprocess
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from datetime import datetime

# Register HEIC opener
pillow_heif.register_heif_opener()

# --- CONSTANTS & LAYOUT ---
CONTENT_WIDTH_IN = 8.5 - (0.5 + 0.5)  # 7.5"
# --- STRICT 1:4 RATIO ---
DESC_COL_IN = CONTENT_WIDTH_IN * (1 / 5.0)   # 1 part = 1.5"
PHOTO_COL_IN = CONTENT_WIDTH_IN * (4 / 5.0)  # 4 parts = 6.0"
ROW_HEIGHT_IN = 4.5  # Two rows per page
PHOTO_MAX_W_IN = PHOTO_COL_IN - 0.20  # Padding
PHOTO_MAX_H_IN = ROW_HEIGHT_IN - 0.20

JOB_MAP = {
    "WM 3155 SSM Parking": "T250014 Walmart 3155 Sault Ste. Marie FY26 Parking Lot Refresh\n446 Great Northern Rd.\nSault Ste. Marie, Ontario, P6B 4Z9",
    "GSDR GSN CC": "T250030 GSDR - Guru Nanak Sewa Community Centre\n1410 Stevenson Rd N\nOshawa, Ontario, L1L 0N6",
}

# --- IMAGE PROCESSING ---
def process_image(uploaded_file, compress=False, quality=70, rotation=0):
    """Open, rotate, and optionally compress an image. Returns a BytesIO stream."""
    try:
        uploaded_file.seek(0)
        img = Image.open(uploaded_file)

        if rotation:
            img = img.rotate(rotation, expand=True)

        if img.mode == 'RGBA':
            bg = Image.new('RGB', img.size, (255, 255, 255))
            bg.paste(img, mask=img.split()[3])
            img = bg
        elif img.mode != 'RGB':
            img = img.convert('RGB')

        img_io = io.BytesIO()
        if compress:
            img.save(img_io, format='JPEG', optimize=True, quality=int(quality))
        else:
            img.save(img_io, format='PNG')
        img_io.seek(0)
        return img_io
    except Exception as e:
        st.error(f"Error processing image '{getattr(uploaded_file, 'name', 'Unknown')}': {e}")
        return None

def get_rotated_preview(uploaded_file, rotation=0):
    """Return a PIL image rotated for preview."""
    try:
        uploaded_file.seek(0)
        img = Image.open(uploaded_file)
        if rotation:
            img = img.rotate(rotation, expand=True)
        if img.mode not in ('RGB', 'RGBA'):
            img = img.convert('RGB')
        return img
    except Exception as e:
        st.error(f"Preview error: {e}")
        return None

# --- DOC HELPERS ---
def add_page_number(paragraph):
    """Adds Page X of Y in footer."""
    paragraph.alignment = WD_ALIGN_PARAGRAPH.RIGHT
    paragraph.add_run("Page ")

    r1 = paragraph.add_run()
    fldChar1 = OxmlElement('w:fldChar')
    fldChar1.set(qn('w:fldCharType'), 'begin')
    r1._r.append(fldChar1)
    instr = OxmlElement('w:instrText')
    instr.set(qn('xml:space'), 'preserve')
    instr.text = "PAGE"
    r1._r.append(instr)
    fldChar2 = OxmlElement('w:fldChar')
    fldChar2.set(qn('w:fldCharType'), 'end')
    r1._r.append(fldChar2)

    paragraph.add_run(" of ")

    r2 = paragraph.add_run()
    fldChar3 = OxmlElement('w:fldChar')
    fldChar3.set(qn('w:fldCharType'), 'begin')
    r2._r.append(fldChar3)
    instr2 = OxmlElement('w:instrText')
    instr2.set(qn('xml:space'), 'preserve')
    instr2.text = "NUMPAGES"
    r2._r.append(instr2)
    fldChar4 = OxmlElement('w:fldChar')
    fldChar4.set(qn('w:fldCharType'), 'end')
    r2._r.append(fldChar4)

def ensure_cell_has_one_empty_paragraph(cell):
    """
    Safely reset a cell without corrupting the DOCX.
    Ensures the cell has exactly one empty paragraph.
    """
    while len(cell.paragraphs) > 1:
        p = cell.paragraphs[-1]
        p._element.getparent().remove(p._element)
    if not cell.paragraphs:
        cell.add_paragraph("")
    else:
        p = cell.paragraphs[0]
        p.text = ""
        p.paragraph_format.space_before = Pt(0)
        p.paragraph_format.space_after = Pt(0)

def set_table_fixed_layout(table):
    """
    Locks table layout so Word doesn't auto-resize columns.
    """
    tbl = table._element
    tblPr = tbl.tblPr
    if tblPr is None:
        tblPr = OxmlElement('w:tblPr')
        tbl.insert(0, tblPr)
    tblLayout = OxmlElement('w:tblLayout')
    tblLayout.set(qn('w:type'), 'fixed')
    tblPr.append(tblLayout)

# --- CORE REPORT GENERATION ---
def generate_report(image_items, job_details, compress_images, quality, progress_bar):
    doc = Document()
    section = doc.sections[0]
    section.left_margin = section.right_margin = Inches(0.5)
    section.top_margin = section.bottom_margin = Inches(0.75)

    # Header
    header = section.header
    header_table = header.add_table(rows=1, cols=2, width=Inches(CONTENT_WIDTH_IN))
    set_table_fixed_layout(header_table)
    header_table.columns[0].width = Inches(2.5)
    header_table.columns[1].width = Inches(CONTENT_WIDTH_IN - 2.5)

    logo_cell = header_table.cell(0, 0)
    ensure_cell_has_one_empty_paragraph(logo_cell)
    if os.path.exists("logo.png"):
        p_logo = logo_cell.paragraphs[0]
        p_logo.add_run().add_picture("logo.png", width=Inches(1.5))
        text_run = p_logo.add_run("\nHKC Construction")
        text_run.font.size = Pt(8)

    info_cell = header_table.cell(0, 1)
    ensure_cell_has_one_empty_paragraph(info_cell)
    p_info = info_cell.paragraphs[0]
    p_info.text = job_details
    p_info.alignment = WD_ALIGN_PARAGRAPH.RIGHT
    info_cell.vertical_alignment = WD_ALIGN_VERTICAL.CENTER

    # Footer
    footer = section.footer
    add_page_number(footer.paragraphs[0])

    # --- LAYOUT FIX: Reverted to the original 2-photo-per-page layout ---
    num_images = len(image_items)
    processed = 0

    for start in range(0, num_images, 2): # Process 2 images per page
        if start > 0:
            doc.add_page_break()

        page_items = image_items[start:start + 2]

        # A 2x2 table: Row 1 is for Photo 1, Row 2 is for Photo 2
        page_table = doc.add_table(rows=2, cols=2)
        page_table.width = Inches(CONTENT_WIDTH_IN)
        set_table_fixed_layout(page_table)
        page_table.style = 'Table Grid'
        page_table.alignment = WD_TABLE_ALIGNMENT.CENTER

        # Set column widths to the strict 1:4 ratio
        page_table.columns[0].width = Inches(DESC_COL_IN)
        page_table.columns[1].width = Inches(PHOTO_COL_IN)
        for r in page_table.rows:
            r.height = Inches(ROW_HEIGHT_IN)
            r.height_rule = WD_ROW_HEIGHT_RULE.EXACTLY
            # Explicitly set cell widths as well for compatibility
            r.cells[0].width = Inches(DESC_COL_IN)
            r.cells[1].width = Inches(PHOTO_COL_IN)

        # Fill the two rows
        for row_idx in range(2):
            # Only proceed if there is an item for this row
            if row_idx < len(page_items):
                item = page_items[row_idx]
                desc_cell = page_table.cell(row_idx, 0)
                photo_cell = page_table.cell(row_idx, 1)

                ensure_cell_has_one_empty_paragraph(desc_cell)
                ensure_cell_has_one_empty_paragraph(photo_cell)

                desc_cell.vertical_alignment = WD_ALIGN_VERTICAL.TOP
                photo_cell.vertical_alignment = WD_ALIGN_VERTICAL.CENTER

                # Description
                p_desc = desc_cell.paragraphs[0]
                run_label = p_desc.add_run("Description:\n")
                run_label.bold = True
                p_desc.add_run(item.get('description', '') or '')

                # Image processing and placement
                rotation = item.get('rotation', 0) or 0
                try:
                    item['file'].seek(0)
                    im = Image.open(item['file'])
                    if rotation:
                        im = im.rotate(rotation, expand=True)
                    w_px, h_px = im.size
                except Exception:
                    w_px, h_px = (1000, 1000)

                img_ratio = (w_px / h_px) if h_px else 1.0
                box_ratio = PHOTO_MAX_W_IN / PHOTO_MAX_H_IN

                img_stream = process_image(
                    item['file'], compress=compress_images, quality=quality, rotation=rotation
                )

                if img_stream:
                    p_photo = photo_cell.paragraphs[0]
                    p_photo.alignment = WD_ALIGN_PARAGRAPH.CENTER
                    if img_ratio >= box_ratio:
                        p_photo.add_run().add_picture(img_stream, width=Inches(PHOTO_MAX_W_IN))
                    else:
                        p_photo.add_run().add_picture(img_stream, height=Inches(PHOTO_MAX_H_IN))

                processed += 1
                if num_images:
                    progress_bar.progress(processed / num_images)

    # Save DOCX
    doc_io = io.BytesIO()
    doc.save(doc_io)
    doc_io.seek(0)

    # Optional PDF conversion
    pdf_io = None
    try:
        with open("temp_report.docx", "wb") as f:
            f.write(doc_io.getvalue())
        subprocess.run(
            ['libreoffice', '--headless', '--convert-to', 'pdf', '--outdir', '.', 'temp_report.docx'],
            check=True, timeout=120
        )
        pdf_path = 'temp_report.pdf'
        if os.path.exists(pdf_path):
            with open(pdf_path, "rb") as f:
                pdf_io = io.BytesIO(f.read())
            os.remove("temp_report.docx")
            os.remove(pdf_path)
    except Exception as e:
        st.warning(f"PDF conversion failed: {e}. You can still download the Word document.")

    return doc_io, pdf_io


# --- STREAMLIT UI ---
st.set_page_config(layout="centered", page_title="HKC Report Generator")

if 'image_dict' not in st.session_state:
    st.session_state.image_dict = {}

def reset_app():
    st.session_state.clear()

if os.path.exists("logo.png"):
    st.image("logo.png", width=200)

st.title("Photo Report Generator")
st.markdown("---")

with st.sidebar:
    st.header("1. Job Information")
    job_key = st.selectbox(
        "Select Job",
        options=list(JOB_MAP.keys()),
        index=0,
        help="This sets the header and the download file name."
    )
    job_details = JOB_MAP[job_key]
    st.text_area("Job Details", job_details, height=125, disabled=True)

    st.header("2. Image Settings")
    compress_images = st.checkbox("Compress Images", value=True)
    quality = st.slider("Compression intensity", min_value=10, max_value=95, value=70, step=5)

    st.header("3. Actions")
    generate_button = st.button("Generate Report", type="primary", use_container_width=True)
    if st.button("Reset All", use_container_width=True):
        reset_app()
        st.rerun()

st.header("Upload, Arrange, Rotate, and Describe Images")
st.info("Arrange photos in the desired order. The report will have 2 photos per page.")
uploaded_files = st.file_uploader(
    "Upload all your photos here",
    accept_multiple_files=True,
    type=["jpg", "jpeg", "png", "heic", "jfif"],
    key="file_uploader"
)

if uploaded_files:
    for file in uploaded_files:
        fid = getattr(file, "file_id", None) or f"{file.name}-{id(file)}"
        if fid not in st.session_state.image_dict:
            current_max_order = max((d.get('order', 0) for d in st.session_state.image_dict.values()), default=0)
            st.session_state.image_dict[fid] = {
                'file': file,
                'description': '',
                'order': current_max_order + 1,
                'rotation': 0
            }

if st.session_state.image_dict:
    items = list(st.session_state.image_dict.items())
    items.sort(key=lambda kv: kv[1].get('order', 0))

    for fid, item_data in items:
        st.markdown("---")
        cols = st.columns([2, 3])

        with cols[0]:
            preview = get_rotated_preview(item_data['file'], item_data.get('rotation', 0))
            if preview is not None:
                st.image(preview, use_column_width=True)
            rcols = st.columns(2)
            if rcols[0].button("↺ Rotate Left", key=f"rotl_{fid}"):
                item_data['rotation'] = (item_data.get('rotation', 0) + 90) % 360
                st.rerun()
            if rcols[1].button("↻ Rotate Right", key=f"rotr_{fid}"):
                item_data['rotation'] = (item_data.get('rotation', 0) - 90) % 360
                st.rerun()

        with cols[1]:
            item_data['order'] = st.number_input(
                "Order", min_value=1, value=item_data.get('order', 1), key=f"order_{fid}"
            )
            item_data['description'] = st.text_area(
                "Description", value=item_data.get('description', ''), key=f"desc_{fid}", height=120
            )

if generate_button:
    if not job_details:
        st.error("Please select a job.")
    elif not st.session_state.image_dict:
        st.error("Please upload at least one image.")
    else:
        sorted_image_items = [v for _, v in sorted(st.session_state.image_dict.items(), key=lambda kv: kv[1].get('order', 0))]
        progress_bar = st.progress(0.0, text="Starting report generation...")
        try:
            with st.spinner("Generating your report..."):
                doc_io, pdf_io = generate_report(
                    sorted_image_items, job_details, compress_images, quality, progress_bar
                )
            progress_bar.success("Report generated successfully!")

            date_str = datetime.now().strftime("%y%m%d")
            file_base = f"{job_key} - {date_str}"

            st.header("Download Your Report")
            c1, c2 = st.columns(2)
            c1.download_button("⬇️ Download Word (.docx)", doc_io, f"{file_base}.docx", use_container_width=True)
            if pdf_io:
                c2.download_button("⬇️ Download PDF (.pdf)", pdf_io, f"{file_base}.pdf", use_container_width=True)
        except Exception as e:
            st.error(f"A critical error occurred: {e}")
            progress_bar.empty()