app.py

import streamlit as st
import pandas as pd
import fitz  # PyMuPDF
import os
import subprocess
import tempfile
import sys
import toml
import shutil
import zipfile
import io

# Ensure we can import from utils if needed
sys.path.append(os.path.dirname(__file__))
from utils import toc_processor
from pdfxmeta import pdfxmeta

st.set_page_config(page_title="PDF Bookmark Splitter", layout="wide")

st.title("PDF Bookmarker & Splitter")

st.markdown("""

**Upload a PDF**, analyze its fonts to find top-level headings, and generate Bookmarks for splitting by chapter.

""")

uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")

# Short guidance shown under the uploader
st.caption("Recommended use: After uploading your PDF, search for the text of a known chapter heading. Once the correct entry is identified in the search results, select the corresponding entry from the drop down, and optionally repeat the step to ensure back matter is split off from the last chapter before running the pipeline.")

if uploaded_file is not None:
    # We need to save the uploaded file to disk for the CLI tools to read it
    # We'll use a permanent temp file for the session so we don't have to re-upload constantly
    # But for cleanliness, we might want to put this in a temp dir too?
    # For now, keeping the input file logic as is (tempfile), but we'll put OUTPUTS in a pure temp dir
    
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_pdf:
        tmp_pdf.write(uploaded_file.getvalue())
        input_pdf_path = tmp_pdf.name

    # --- State Management & Reset ---
    # Check if a new file is uploaded
    file_id = f"{uploaded_file.name}_{uploaded_file.size}" # Robust proxy for ID
    if 'current_file_id' not in st.session_state:
        st.session_state['current_file_id'] = None
        
    if st.session_state['current_file_id'] != file_id:
        # NEW FILE DETECTED: Reset Pipeline State
        keys_to_reset = ['final_pdf_bytes', 'final_zip_bytes', 'final_zip_name', 'search_matches', 'font_name', 'font_size']
        for k in keys_to_reset:
            if k in st.session_state:
                del st.session_state[k]
        st.session_state['current_file_id'] = file_id
        # st.toast(f"New file loaded: {uploaded_file.name}. State cleared.")

    st.success(f"Loaded: {uploaded_file.name}")

    # --- Data Source Selection ---
    st.header("1. Source Selection")
    source_mode = st.radio("Where should the bookmarks come from?", 
             ["Scan & Generate (Create New)", "Use Existing Bookmarks (Modify)"],
             help="Choose 'Scan & Generate' to build new bookmarks from fonts. Choose 'Use Existing' to tidy up bookmarks already in the file.")

    # --- Analysis Section (Only for Generate) ---
    if source_mode == "Scan & Generate (Create New)":
        st.header("2. Analyze Fonts")
        
        if 'font_name' not in st.session_state:
            st.session_state['font_name'] = ''
        if 'font_size' not in st.session_state:
            st.session_state['font_size'] = 18.0
            
        tab1, tab2 = st.tabs(["Scan for Large Fonts", "Search by Text"])
        
        with tab1:
            if st.button("Find Header Candidates"):
                with st.spinner("Scanning PDF for large fonts..."):
                    doc = fitz.open(input_pdf_path)
                    candidates = []
                    for page in doc[:50]:
                        text_page = page.get_text("dict")
                        for block in text_page["blocks"]:
                            for line in block.get("lines", []):
                                for span in line["spans"]:
                                    text = span["text"].strip()
                                    if len(text) > 3:
                                        candidates.append({
                                            "Text": text[:50],
                                            "Font": span["font"],
                                            "Size": round(span["size"], 2),
                                            "Page": page.number + 1
                                        })
                    doc.close()
                    if candidates:
                        df = pd.DataFrame(candidates)
                        summary = df.groupby(['Font', 'Size']).size().reset_index(name='Count')
                        summary = summary.sort_values(by=['Size', 'Count'], ascending=[False, False]).head(20)
                        st.session_state['scan_results'] = summary
                    else:
                        st.warning("No text found.")
            
            if 'scan_results' in st.session_state:
                st.write("### Top Large Fonts Found")
                st.dataframe(st.session_state['scan_results'], use_container_width=True)
                
                def update_from_scan():
                    val = st.session_state.scan_selector
                    if val:
                        f_name = val.split(" (")[0]
                        f_size = float(val.split("(")[1].replace("pt)", ""))
                        st.session_state['font_name'] = f_name
                        st.session_state['font_size'] = f_size

                options = st.session_state['scan_results'].apply(lambda x: f"{x['Font']} ({x['Size']}pt)", axis=1)
                st.selectbox("Select extraction font:", options, key='scan_selector', on_change=update_from_scan, index=None, placeholder="Choose a font...")

        with tab2:
            search_query = st.text_input("Enter text to find (e.g., 'Chapter 1')", "")
            
            c1, c2 = st.columns([1, 3])
            with c1:
                do_search = st.button("Search Text")
            with c2:
                is_case_sensitive = st.checkbox("Case Sensitive", value=False)

            if do_search:
                with st.spinner(f"Searching for '{search_query}'..."):
                    # Use the robust pdfxmeta library
                    try:
                        doc = fitz.open(input_pdf_path)
                        # pdfxmeta expects a regex pattern, so we escape the query to be safe
                        import re
                        safe_pattern = re.escape(search_query)
                        
                        # extract_meta returns a list of dicts (spans)
                        results = pdfxmeta.extract_meta(doc, safe_pattern, ign_case=(not is_case_sensitive))
                        doc.close()
                        
                        matches = []
                        for res in results:
                            matches.append({
                                "Text": res.get("text", "").strip(),
                                "Font": res.get("font", ""),
                                "Size": round(res.get("size", 0), 2),
                                "Page": res.get("page_index", 0)
                            })
                            # Limit for display safety
                            if len(matches) > 50: break
                            
                        if matches:
                            st.session_state['search_matches'] = pd.DataFrame(matches)
                        else:
                            st.warning("No matches found.")
                            
                    except Exception as e:
                        st.error(f"Search failed: {e}")
            
            if 'search_matches' in st.session_state:
                st.write(f"### Found Matches")
                st.dataframe(st.session_state['search_matches'], use_container_width=True)
                
                def update_from_search():
                    val = st.session_state.search_selector
                    if val:
                        parts = val.split(" (")
                        f_name = parts[0]
                        f_size = float(parts[1].split("pt)")[0])
                        st.session_state['font_name'] = f_name
                        st.session_state['font_size'] = f_size

                options = st.session_state['search_matches'].apply(lambda x: f"{x['Font']} ({x['Size']}pt) - Pg {x['Page']}", axis=1)
                st.selectbox("Select font from match:", options, key='search_selector', on_change=update_from_search, index=None, placeholder="Choose a match...")

        # --- Configuration (Only for Generate) ---
        st.header("3. Configure Recipe")
        col1, col2 = st.columns(2)
        with col1:
            font_name_input = st.text_input("Font Name", key='font_name')
        with col2:
            font_size_input = st.number_input("Font Size", key='font_size')
        
        greedy = st.checkbox("Greedy Match (Merge multiline specs)", value=True)
        
        # --- Back Matter Configuration ---
        with st.expander("Back Matter Configuration (Optional)", expanded=False):
            st.markdown("Identify where the **Back Matter** (Index, Glossary, etc.) starts to split it into a separate `999_Back_matter.pdf`.")
            
            # Independent Search for Back Matter
            bm_query = st.text_input("Find Back Matter start (e.g., 'Index')", key="bm_search_query")
            
            c_bm1, c_bm2 = st.columns([1, 3])
            with c_bm1:
                 do_bm_search = st.button("Search Back Matter")
            with c_bm2:
                 bm_case_sensitive = st.checkbox("Case Sensitive", key="bm_sens", value=False)
                 
            if do_bm_search:
                with st.spinner("Searching..."):
                    try:
                        doc = fitz.open(input_pdf_path)
                        import re
                        safe_pattern = re.escape(bm_query)
                        results = pdfxmeta.extract_meta(doc, safe_pattern, ign_case=(not bm_case_sensitive))
                        doc.close()
                        
                        bm_matches = []
                        for res in results:
                            bm_matches.append({
                                "Text": res.get("text", "").strip(), 
                                "Page": res.get("page_index", 0) # Display raw (already 1-based from pdfxmeta)
                            })
                            if len(bm_matches) > 50: break
                        
                        if bm_matches:
                            st.session_state['bm_matches'] = pd.DataFrame(bm_matches)
                        else:
                            st.warning("No matches found.")
                    except Exception as e:
                        st.error(f"Search failed: {e}")

            if 'bm_matches' in st.session_state:
                st.dataframe(st.session_state['bm_matches'], use_container_width=True)
                
                def update_bm_page():
                    val = st.session_state.bm_selector
                    if val:
                        # Value format: "Page X - Text..."
                        page_num = int(val.split(" -")[0].replace("Page ", ""))
                        st.session_state['back_matter_page'] = page_num
                
                bm_options = st.session_state['bm_matches'].apply(lambda x: f"Page {x['Page']} - {x['Text'][:30]}...", axis=1)
                st.selectbox("Select Start Page:", bm_options, key='bm_selector', on_change=update_bm_page, index=None, placeholder="Select start page...")

            # Manual Override
            # Update session state when this input changes
            def update_manual_bm():
                st.session_state['back_matter_page'] = st.session_state.back_matter_page_manual
                
            st.number_input("Or manually set Start Page:", min_value=0, value=st.session_state.get('back_matter_page', 0), key='back_matter_page_manual', on_change=update_manual_bm)

    else:
        # Existing Mode
        st.info("Using existing bookmarks. They will be cleaned, numbered, and used for splitting/downloading.")

    # --- Generation ---
    st.header("4. Process & Generate")
    
    if st.button("Run Pipeline"):
        # Validate inputs if generating
        if source_mode == "Scan & Generate (Create New)" and not st.session_state.get('font_name'):
            st.error("Please specify a font name for extraction.")
        else:
            with st.status("Running pipeline tasks...", expanded=True) as status:
                # Use a temporary directory for all intermediate files
                with tempfile.TemporaryDirectory() as temp_dir:
                    status.write(f"Created temp workspace: {temp_dir}")
                    
                    # Paths
                    recipe_path = os.path.join(temp_dir, "recipe.toml")
                    raw_toc_path = os.path.join(temp_dir, "raw.toc") # pdftocgen output
                    clean_toc_path = os.path.join(temp_dir, "clean.toc") # modify_toc output
                    output_pdf_path = os.path.join(temp_dir, "final.pdf")
                    
                    raw_toc_content = ""

                    if source_mode == "Scan & Generate (Create New)":
                        # 1. Create Recipe
                        recipe_data = {
                            "heading": [{
                                "level": 1,
                                "greedy": greedy,
                                "font": {
                                    "name": st.session_state['font_name'],
                                    "size": st.session_state['font_size'],
                                    "size_tolerance": 0.1
                                }
                            }]
                        }
                        with open(recipe_path, "w") as f:
                            toml.dump(recipe_data, f)
                        status.write("✅ Recipe created")
                        
                        # 2. Run pdftocgen -> raw.toc
                        status.write("Running pdftocgen (Scanning)...")
                        cmd1 = [sys.executable, "-m", "pdftocgen", "-r", recipe_path, input_pdf_path]
                        process = subprocess.run(cmd1, capture_output=True, text=True, encoding='utf-8')
                        if process.returncode != 0:
                            st.error(f"pdftocgen failed: {process.stderr}")
                            st.stop()
                        raw_toc_content = process.stdout
                        status.write("✅ Headers extracted")

                    else:
                        # Existing Bookmarks
                        status.write("Extracting existing bookmarks...")
                        # Run pdftocio in extract mode
                        cmd1 = [sys.executable, "-m", "pdftocio", input_pdf_path]
                        process = subprocess.run(cmd1, capture_output=True, text=True, encoding='utf-8')
                        if process.returncode != 0:
                            st.error(f"pdftocio failed: {process.stderr}")
                            st.stop()
                        raw_toc_content = process.stdout
                        if not raw_toc_content.strip():
                            st.warning("No existing bookmarks found!")
                            st.stop()
                        status.write("✅ Existing bookmarks imported")
                    
                    # 3. Clean Content (Using centralized utility)
                    status.write("Cleaning and merging bookmarks...")
                    cleaned_toc_content = toc_processor.process_toc(raw_toc_content)
                    
                    with open(clean_toc_path, "w", encoding='utf-8') as f:
                        f.write(cleaned_toc_content)
                    status.write("✅ Bookmarks formatted (Double-splits fixed)")
                    
                    # 4. Write PDF
                    status.write("Writing to PDF...")
                    cmd3 = [sys.executable, "-m", "pdftocio", "-t", clean_toc_path, "-o", output_pdf_path, input_pdf_path]
                    process = subprocess.run(cmd3, capture_output=True, text=True)
                    if process.returncode != 0:
                        st.error(f"pdftocio failed: {process.stderr}")
                        st.stop()
                    status.write("✅ PDF saved")
                    
                    # 5. Read Result for Download
                    with open(output_pdf_path, "rb") as f:
                        st.session_state['final_pdf_bytes'] = f.read()
                    
                    # 6. Split & Zip (The Feature)
                    # Use a temp file for the zip to avoid memory issues
                    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp_zip:
                        tmp_zip_path = tmp_zip.name
                    
                    try:
                        # Pass back_matter_page if it exists and is valid
                        bm_page = st.session_state.get('back_matter_page', 0)
                        if bm_page == 0: bm_page = None
                        
                        toc_processor.generate_chapter_splits(output_pdf_path, tmp_zip_path, back_matter_start_page=bm_page)
                        
                        with open(tmp_zip_path, "rb") as f:
                            st.session_state['final_zip_bytes'] = f.read()
                            
                        base_name = os.path.splitext(uploaded_file.name)[0]
                        st.session_state['final_zip_name'] = f"{base_name}_chapters.zip"
                        
                    except Exception as e:
                        st.error(f"Error generating zip: {e}")
                    finally:
                        if os.path.exists(tmp_zip_path):
                            os.unlink(tmp_zip_path)

    # --- Persistent Download Area ---
    if 'final_pdf_bytes' in st.session_state:
        st.success("Pipeline completed successfully!")
        st.write("### Downloads")
        
        c_dl1, c_dl2 = st.columns(2)
        with c_dl1:
            st.download_button(
                label="Download Bookmarked PDF",
                data=st.session_state['final_pdf_bytes'],
                file_name="bookmarked_doc.pdf",
                mime="application/pdf",
                key="dl_pdf_btn"
            )
        
        with c_dl2:
            if 'final_zip_bytes' in st.session_state:
                st.download_button(
                    label=f"Download ZIP ({st.session_state['final_zip_name']})",
                    data=st.session_state['final_zip_bytes'],
                    file_name=st.session_state['final_zip_name'],
                    mime="application/zip",
                    key="dl_zip_btn"
                )

    st.markdown("---")
    st.markdown("""

    <div style="text-align: center; color: #666; font-size: 0.8em;">

        Based on <a href="https://github.com/Krasjet/pdf.tocgen" target="_blank">pdf.tocgen</a> by krasjet. <br>

        Enhanced with UI, Chapter Splitting, and Metadata Search. Licensed under AGPL-3.0.

    </div>

    """, unsafe_allow_html=True)