import streamlit as st
import fitz
from PIL import Image
import io

st.header("Line Art Data Annotation")
uploaded_pdf = st.sidebar.file_uploader("Upload a PDF", type=["pdf"])

if uploaded_pdf:
    data = uploaded_pdf.read()
    doc = fitz.open(stream=data, filetype="pdf")
    
    # Initialize page index in session state
    if "page_idx" not in st.session_state:
        st.session_state.page_idx = 0

    total_pages = doc.page_count
    page_idx = st.session_state.page_idx % total_pages
    # Navigation buttons (placed above the image)
    col_prev, col_caption, col_next = st.columns([1, 8, 1])
    with col_prev:
        if st.button("<"):
            st.session_state.page_idx = (page_idx - 1) % total_pages
            st.rerun()
    with col_caption:
        st.markdown(f"<center>Page {page_idx}/{total_pages - 1}</center>", unsafe_allow_html=True)
    with col_next:
        if st.button("\>"):
            st.session_state.page_idx = (page_idx + 1) % total_pages
            st.rerun()

    # Render image after controls
    page = doc.load_page(page_idx)
    pix = page.get_pixmap(dpi=200)
    img = Image.open(io.BytesIO(pix.tobytes("png")))
    st.image(img)