import streamlit as st
from PIL import Image, ImageEnhance
from io import BytesIO
import zipfile
import os
from fpdf import FPDF
import hashlib

st.set_page_config(page_title="Smart Doc Scanner", layout="wide")

def resize_image(img, max_height=600):
    if img.height > max_height:
        ratio = max_height / img.height
        return img.resize((int(img.width * ratio), max_height), Image.Resampling.LANCZOS)
    return img

def crop_sliders(img, key):
    w, h = img.size
    state = st.session_state
    for name, val in zip(["left", "top", "right", "bottom"], [0, 0, w, h]):
        state.setdefault(f"{key}_{name}", val)

    col1, col2 = st.columns(2)
    with col1:
        st.session_state[f"{key}_left"] = st.slider("Crop Left", 0, w - 10, st.session_state[f"{key}_left"], key=f"{key}_left_slider")
        st.session_state[f"{key}_right"] = st.slider("Crop Right", st.session_state[f"{key}_left"] + 10, w, st.session_state[f"{key}_right"], key=f"{key}_right_slider")
    with col2:
        st.session_state[f"{key}_top"] = st.slider("Crop Top", 0, h - 10, st.session_state[f"{key}_top"], key=f"{key}_top_slider")
        st.session_state[f"{key}_bottom"] = st.slider("Crop Bottom", st.session_state[f"{key}_top"] + 10, h, st.session_state[f"{key}_bottom"], key=f"{key}_bottom_slider")

    return (
        st.session_state[f"{key}_left"],
        st.session_state[f"{key}_top"],
        st.session_state[f"{key}_right"],
        st.session_state[f"{key}_bottom"]
    )

st.title("📄 Smart ID & Document Scanner")
uploaded_files = st.file_uploader("📤 Upload Images", type=["png", "jpg", "jpeg"], accept_multiple_files=True)

if uploaded_files:
    mode = st.radio("Mode", ["Normal Mode", "Advanced Mode (Edit Each)"])
    processed_images = []
    file_names = []

    for file in uploaded_files:
        file_bytes = file.getvalue()
        hash_key = hashlib.md5(file_bytes).hexdigest()
        key = f"img_{hash_key}"

        st.markdown(f"---\n### 🖼️ {file.name}")
        original = Image.open(BytesIO(file_bytes)).convert("RGB")

        # Init session state
        if f"original_{key}" not in st.session_state:
            st.session_state[f"original_{key}"] = original
            st.session_state[f"enhance_{key}"] = False
            st.session_state[f"filename_{key}"] = file.name.rsplit(".", 1)[0]
            w, h = original.size
            st.session_state[f"{key}_left"] = 0
            st.session_state[f"{key}_top"] = 0
            st.session_state[f"{key}_right"] = w
            st.session_state[f"{key}_bottom"] = h
            st.session_state[f"angle_{key}"] = 0

        # ✅ NORMAL MODE HANDLING
        if mode == "Normal Mode":
            st.image(resize_image(original), use_container_width=True)
            processed_images.append(original)
            file_names.append(st.session_state[f"filename_{key}"])
            continue

        # ADVANCED MODE
        col1, col2 = st.columns([1, 1.5])
        with col1:
            if st.button("🔁 Reset", key=f"reset_{key}"):
                st.session_state[f"angle_{key}"] = 0
                st.session_state[f"enhance_{key}"] = False
                w, h = original.size
                st.session_state[f"{key}_left"] = 0
                st.session_state[f"{key}_top"] = 0
                st.session_state[f"{key}_right"] = w
                st.session_state[f"{key}_bottom"] = h

            angle = st.slider("🌀 Rotate", -180, 180, value=st.session_state[f"angle_{key}"], key=f"angle_{key}_slider")
            st.session_state[f"angle_{key}"] = angle

            st.checkbox("✨ Enhance", key=f"enhance_{key}")
            left, top, right, bottom = crop_sliders(original, key)
            st.text_input("📁 Rename this file", value=st.session_state[f"filename_{key}"], key=f"filename_{key}_input")

        with col2:
            angle = st.session_state[f"angle_{key}"]
            enhance = st.session_state[f"enhance_{key}"]

            img = st.session_state[f"original_{key}"].crop((left, top, right, bottom))
            img = resize_image(img)

            if enhance:
                img = ImageEnhance.Contrast(img).enhance(1.5)
                img = ImageEnhance.Sharpness(img).enhance(2.0)

            if angle != 0:
                img = img.rotate(angle, expand=True)

            st.image(img, caption=f"🎯 Live Preview ({angle}°)", use_container_width=True)

            if st.button("✅ Save This Image", key=f"save_{key}"):
                st.session_state[f"final_{key}"] = img.copy()
                st.success("Saved for export!")

        final_img = st.session_state.get(f"final_{key}", original)
        filename = st.session_state.get(f"filename_{key}", file.name.rsplit(".", 1)[0])
        processed_images.append(final_img)
        file_names.append(filename)

    # --- EXPORT SECTION ---
    st.markdown("### 🧾 Export Options")

    # ZIP layout
    zip_col1, zip_col2 = st.columns([1, 2])
    with zip_col1:
        if st.button("📦 Generate ZIP"):
            zip_buffer = BytesIO()
            with zipfile.ZipFile(zip_buffer, "w") as zipf:
                for name, img in zip(file_names, processed_images):
                    buffer = BytesIO()
                    img.save(buffer, format="PNG")
                    zipf.writestr(f"{name}.png", buffer.getvalue())
            zip_buffer.seek(0)
            st.session_state["zip_file"] = zip_buffer
    with zip_col2:
        if "zip_file" in st.session_state:
            st.download_button("⬇️ Download ZIP", st.session_state["zip_file"], "images.zip", mime="application/zip", key="zip_dl")

    # PDF layout
    pdf_col1, pdf_col2 = st.columns([1, 2])
    with pdf_col1:
        if st.button("📄 Generate PDF"):
            pdf = FPDF()
            tmp_paths = []
            for name, img in zip(file_names, processed_images):
                path = f"/tmp/{name}.jpg"
                img.save(path, "JPEG")
                pdf.add_page()
                pdf.image(path, x=10, y=10, w=190)
                tmp_paths.append(path)
            out_path = "/tmp/final.pdf"
            pdf.output(out_path)
            with open(out_path, "rb") as f:
                st.session_state["pdf_file"] = f.read()
            for p in tmp_paths + [out_path]:
                os.remove(p)
    with pdf_col2:
        if "pdf_file" in st.session_state:
            st.download_button("⬇️ Download PDF", st.session_state["pdf_file"], "scanned.pdf", mime="application/pdf", key="pdf_dl")