app.py

import streamlit as st
import pandas as pd
import json
import time # For simulating delays

# --- Configuration and Initial Setup ---
st.set_page_config(layout="wide", page_title="Dummy URS Doc Processor UI Test")
st.title("Dummy URS Doc Processor - UI Test Version")

# Initialize session states
if "show_step2" not in st.session_state:
    st.session_state.show_step2 = False
if "show_step3" not in st.session_state:
    st.session_state.show_step3 = False
if "protocol_generated" not in st.session_state:
    st.session_state.protocol_generated = False
if "extracted_data" not in st.session_state:
    st.session_state.extracted_data = None
if "dataframe" not in st.session_state:
    st.session_state.dataframe = None
if "processing" not in st.session_state:
    st.session_state.processing = False
if "generated_pdf" not in st.session_state:
    st.session_state.generated_pdf = None
if "generated_pdf_name" not in st.session_state:
    st.session_state.generated_pdf_name = None

# Dummy LLM setup (no actual LLM calls)
@st.cache_resource
def get_llms():
    return None, None # No LLMs needed for UI test

# --------------------- Dummy Backend Functions ---------------------

def dummy_run_rag_pipeline(pdf_file, template_file):
    """
    Simulates the RAG pipeline for UI testing.
    Returns dummy extracted data.
    """
    st.toast("Simulating RAG Pipeline...")
    time.sleep(2) # Simulate processing time

    # Dummy extracted data
    dummy_data = {
        "ProjectName": "Dummy Project UI",
        "DocumentNumber": "DUMMY-001",
        "SystemName": "Test System Beta",
        "Version": "1.0.0",
        "Author": "UI Tester",
        "Date": "2025-07-21",
        "FunctionalRequirement_1": "The system shall display data responsively on mobile devices.",
        "FunctionalRequirement_2": "The system shall allow editing of extracted values.",
        "Utility#1": "Power Supply (240V)",
        "Utility#2": "Ethernet Connection",
        "SoftwareRequirements": "Python 3.9+, Streamlit 1.x",
        "HardwareRequirements": "Any modern smartphone or tablet"
    }
    st.toast("RAG Pipeline simulation complete!")
    return dummy_data

def dummy_json_to_dataframe(json_data):
    """
    Converts dummy JSON data to DataFrame for editing.
    """
    if not json_data:
        return pd.DataFrame(columns=['Placeholder', 'Value'])

    data_list = []
    for key, value in json_data.items():
        data_list.append((str(key), str(value) if value is not None else ""))
    
    df = pd.DataFrame(data_list, columns=['Placeholder', 'Value'])
    return df

def create_dummy_pdf():
    """Generates a dummy PDF byte stream for download."""
    from reportlab.lib.pagesizes import letter
    from reportlab.pdfgen import canvas
    from io import BytesIO

    buffer = BytesIO()
    p = canvas.Canvas(buffer, pagesize=letter)
    p.drawString(100, 750, "This is a Dummy PDF for UI Testing!")
    p.drawString(100, 730, "It contains no real data from your URS or protocol.")
    p.drawString(100, 710, "Generated on: " + time.ctime())
    p.showPage()
    p.save()
    buffer.seek(0)
    return buffer.getvalue()

# --- Status & Progress Tracking ---
# This section remains largely the same to simulate your original progress logic
progress = 0
status_text = "🟦 Waiting for input files"

if st.session_state.processing:
    progress = 20
    status_text = "Extracting values from uploaded files (simulated)..."

if st.session_state.show_step2:
    progress = 40
    status_text = "Input form is being generated (simulated)..."

if st.session_state.show_step3:
    progress = 60
    status_text = "Saving changes to input form (simulated)..."

if st.session_state.protocol_generated:
    progress = 100
    status_text = "Final protocol document is being generated (simulated)..."

st.markdown(f"### Status: {status_text}")
st.progress(progress / 100.0)

# --- Layout: Horizontal columns for each step ---
st.write("---") # Horizontal line for separation
col1, col2, col3 = st.columns(3)

# --- STEP 1 ---
with col1:
    st.markdown("### 🔹 Step 1")
    st.write("📄 **Upload input URS (PDF)**")
    urs_pdf = st.file_uploader("Upload PDF", type=["pdf"], key="pdf_upload_step1")

    st.write("📊 **Upload template (XLSX)**")
    template_xlsx = st.file_uploader("Upload XLSX", type=["xlsx"], key="xlsx_upload_step1")

    if st.button("➡️ Generate Input Form", key="generate_form_button"):
        if urs_pdf and template_xlsx:
            st.session_state.processing = True
            st.rerun()
        else:
            st.warning("Please upload both PDF and XLSX files for the simulation.")

# Process RAG pipeline if processing is True
if st.session_state.processing:
    with st.spinner("🔄 Running RAG Pipeline (Simulated)..."):
        # Call the dummy RAG pipeline
        result = dummy_run_rag_pipeline(urs_pdf, template_xlsx)
        
        if result:
            st.session_state.extracted_data = result
            st.session_state.dataframe = dummy_json_to_dataframe(result)
            st.session_state.show_step2 = True
            st.session_state.processing = False
            st.success("✅ RAG Pipeline simulation completed!")
            st.rerun()
        else:
            st.error("❌ Failed to simulate document processing. Try again.")
            st.session_state.processing = False
            st.rerun()

st.write("---") # Horizontal line for separation

# --- STEP 2 ---
with col2:
    if st.session_state.show_step2:
        st.markdown("### 🔹 Step 2")
        
        if st.session_state.dataframe is not None:
            st.markdown("✏️ **Edit Input Form (DataFrame)**")
            
            # Display editable dataframe
            edited_df = st.data_editor(
                st.session_state.dataframe,
                num_rows="dynamic",
                key="data_editor",
                height=400 # Fixed height for better mobile consistency
            )
            
            # Update session state with edited data (for dummy purposes)
            st.session_state.dataframe = edited_df
            
            if st.button("💾 Save Changes", key="save_changes_button"):
                # Simulate saving changes
                st.session_state.extracted_data = {
                    row['Placeholder']: row['Value'] 
                    for _, row in edited_df.iterrows() 
                    if pd.notna(row['Value']) and row['Value'] != ''
                }
                st.success("✅ Changes saved (simulated)!")
                st.session_state.show_step3 = True
                st.rerun()
    else:
        st.markdown("### 🔹 Step 2")
        st.info("Complete Step 1 to unlock.")

st.write("---") # Horizontal line for separation

# --- STEP 3 ---
with col3:
    if st.session_state.show_step3:
        st.markdown("### 🔹 Step 3")
        st.write("📤 **Upload Protocol Template (DOCX)**")
        output_docx = st.file_uploader("Upload DOCX", type=["docx"], key="docx_upload_step3")

        # Dummy progress bar for protocol generation steps
        protocol_status_text = st.empty()
        protocol_progress_bar_inner = st.progress(0.0)

        if st.button("📝 Generate Protocol Document", key="generate_protocol_button"):
            if output_docx is None:
                st.warning("Please upload a DOCX template file for the simulation.")
            else:
                protocol_status_text.text("💾 Saving completed DOCX (simulated)...")
                protocol_progress_bar_inner.progress(0.5)
                time.sleep(1) # Simulate delay

                protocol_status_text.text("⏳ Converting DOCX to PDF (simulated)...")
                protocol_progress_bar_inner.progress(0.7)
                time.sleep(1) # Simulate delay

                # Generate a dummy PDF
                st.session_state.generated_pdf = create_dummy_pdf()
                st.session_state.generated_pdf_name = "dummy_final_protocol.pdf"
                st.session_state.protocol_generated = True
                
                protocol_status_text.text("✅ Final protocol document generated (simulated)!")
                protocol_progress_bar_inner.progress(1.0)
                st.success("✅ Protocol document generated! You can now view or download the dummy PDF.")
                st.rerun()

        if st.session_state.protocol_generated and hasattr(st.session_state, 'generated_pdf'):
            st.markdown("#### 📥 View / Download")
            col_a, col_b = st.columns(2) # Removed 'Save' column
            with col_a:
                # The 'View PDF' button essentially just triggers a download with preview capabilities
                st.download_button(
                    label="👁️ View PDF",
                    data=st.session_state.generated_pdf,
                    file_name=st.session_state.generated_pdf_name,
                    mime="application/pdf",
                    key="view_pdf_button"
                )
            with col_b:
                st.download_button(
                    label="⬇️ Download PDF",
                    data=st.session_state.generated_pdf,
                    file_name=st.session_state.generated_pdf_name,
                    mime="application/pdf",
                    key="download_pdf_button"
                )
    else:
        st.markdown("### 🔹 Step 3")
        st.info("Complete Step 2 to unlock.")

st.write("---") # Final horizontal line