import streamlit as st
import os
import logging
import warnings
from io import BytesIO
import tempfile
from main import (
    load_workbook_with_possible_passwords,
    copy_excel_file,
    remove_all_protection_tags,
    setup_logging
)

# Configure page
st.set_page_config(
    page_title="Office Protection Remover",
    page_icon="🔓",
    layout="wide"
)

# Setup logging
setup_logging()
warnings.filterwarnings('ignore', category=UserWarning, module='openpyxl.reader.workbook')

def save_uploaded_file(uploaded_file):
    """Save an uploaded file to a temporary location"""
    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_file:
        tmp_file.write(uploaded_file.getvalue())
        return tmp_file.name

# Main UI
st.title("🔓 Office File Protection Remover")
st.write("Remove protection from Excel and Word files easily")

# Sidebar
with st.sidebar:
    st.header("Settings")
    file_type = st.radio(
        "Choose file type:",
        ("Excel", "Word"),
        help="Select the type of files you want to process"
    )

# Main content area
if file_type == "Excel":
    st.header("Excel File Processing")
    
    col1, col2 = st.columns(2)
    
    with col1:
        uploaded_files = st.file_uploader(
            "Upload Excel files", 
            type=["xlsx", "xlsm"],
            accept_multiple_files=True,
            help="You can upload multiple Excel files"
        )
    
    with col2:
        password_option = st.radio(
            "Password Option:",
            ("No Password", "Single Password", "Password File")
        )
        
        passwords = []
        if password_option == "Single Password":
            password = st.text_input("Enter password:", type="password")
            if password:
                passwords = [password]
        elif password_option == "Password File":
            password_file = st.file_uploader("Upload password file", type=["txt"])
            if password_file:
                content = password_file.getvalue().decode()
                passwords = [line.strip() for line in content.splitlines() if line.strip()]
                st.info(f"Loaded {len(passwords)} passwords")

else:  # Word
    st.header("Word Document Processing")
    uploaded_files = st.file_uploader(
        "Upload Word files", 
        type=["docx", "docm"],
        accept_multiple_files=True,
        help="You can upload multiple Word files"
    )

if uploaded_files:
    if st.button(f"Process {file_type} Files", type="primary"):
        progress_bar = st.progress(0)
        status_text = st.empty()
        
        for idx, uploaded_file in enumerate(uploaded_files):
            try:
                # Create a container for each file
                with st.expander(f"Processing {uploaded_file.name}", expanded=True):
                    st.write(f"📝 Processing {uploaded_file.name}...")
                    
                    # Save uploaded file temporarily
                    temp_input_path = save_uploaded_file(uploaded_file)
                    temp_output_path = os.path.join(
                        os.path.dirname(temp_input_path),
                        f"processed_{os.path.basename(temp_input_path)}"
                    )
                    
                    # Process based on file type
                    if file_type == "Excel":
                        copy_excel_file(temp_input_path, temp_output_path, passwords)
                    else:  # Word
                        remove_all_protection_tags(temp_input_path, temp_output_path)
                    
                    # Provide download button
                    with open(temp_output_path, "rb") as f:
                        processed_file = f.read()
                    
                    mime_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" if file_type == "Excel" else "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
                    
                    st.download_button(
                        label=f"⬇️ Download processed file",
                        data=processed_file,
                        file_name=f"processed_{uploaded_file.name}",
                        mime=mime_type
                    )
                    
                    # Cleanup temporary files
                    os.unlink(temp_input_path)
                    os.unlink(temp_output_path)
                    
                    st.success("✅ Processing complete!")
                
                # Update progress
                progress = (idx + 1) / len(uploaded_files)
                progress_bar.progress(progress)
                status_text.text(f"Processed {idx + 1} of {len(uploaded_files)} files")
                
            except Exception as e:
                st.error(f"❌ Error processing {uploaded_file.name}: {str(e)}")
        
        progress_bar.empty()
        status_text.text("✨ All processing complete!")

# Footer
st.sidebar.markdown("---")
st.sidebar.markdown("### Instructions")
st.sidebar.markdown("""
1. Select file type (Excel or Word)
2. Upload your files by dragging them or clicking the upload area
3. For Excel files:
   - Choose password option if needed
   - Enter password or upload password file
4. Click the Process button
5. Download each processed file
""")

st.sidebar.markdown("---")
st.sidebar.markdown("### About")
st.sidebar.markdown("""
This tool helps you remove protection from:
- Excel workbooks (.xlsx, .xlsm)
- Word documents (.docx, .docm)

No files are stored on the server - all processing happens in your browser!
""")