import streamlit as st
from streamlit import components
import os
import logging
import warnings
import shutil
from io import BytesIO
import tempfile
import zipfile
from pathlib import Path
import glob
import platform
import subprocess
from main import (
    load_workbook_with_possible_passwords,
    copy_excel_file,
    remove_all_protection_tags,
    setup_logging
)


def is_running_locally():
    """Check if the app is running locally or in cloud environment"""
    return os.environ.get('STREAMLIT_SERVER_ADDRESS', '').startswith('localhost')

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

# Custom CSS for wider buttons
st.markdown("""
    <style>
        .stButton > button {
            width: 100%;
            min-width: 200px;
        }
        /* Make text areas wider */
        .stTextArea textarea {
            width: 100%;
            min-width: 200px;
        }
        /* Make error section buttons full width */
        [data-testid="stHorizontalBlock"] .stButton button {
            width: 100%;
            min-width: 200px;
        }
    </style>
""", unsafe_allow_html=True)

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

def add_to_error_log(error_dict, filepath, error):
    """Add error to the error dictionary with file path as key"""
    error_dict[filepath] = str(error)

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

def get_file_path(title="Select File", file_types=[("Text Files", "*.txt")]):
    """Get file path using native file dialog when possible"""
    try:
        if platform.system() == "Windows":
            import tkinter as tk
            from tkinter import filedialog
            root = tk.Tk()
            root.withdraw()  # Hide the main window
            root.wm_attributes('-topmost', 1)  # Bring dialog to front
            path = filedialog.askopenfilename(title=title, filetypes=file_types)
            root.destroy()  # Explicitly destroy the tkinter instance
            return path if path else None
        elif platform.system() == "Linux":
            try:
                result = subprocess.run(
                    ['zenity', '--file-selection', '--title', title],
                    capture_output=True,
                    text=True
                )
                return result.stdout.strip() if result.returncode == 0 else None
            except FileNotFoundError:
                st.error("Zenity is not installed. Please install it for file dialog support.")
                return None
        elif platform.system() == "Darwin":  # macOS
            try:
                result = subprocess.run(
                    ['osascript', '-e', f'choose file with prompt "{title}"'],
                    capture_output=True,
                    text=True
                )
                return result.stdout.strip() if result.returncode == 0 else None
            except FileNotFoundError:
                st.error("AppleScript is not available for file dialog support.")
                return None
    except Exception as e:
        st.error(f"Error opening file dialog: {str(e)}")
        return None
    return None

def get_directory_path(title="Select Directory"):
    """Get directory path using native file dialog when possible"""
    try:
        if platform.system() == "Windows":
            import tkinter as tk
            from tkinter import filedialog
            root = tk.Tk()
            root.withdraw()  # Hide the main window
            root.wm_attributes('-topmost', 1)  # Bring dialog to front
            path = filedialog.askdirectory(title=title)
            root.destroy()  # Explicitly destroy the tkinter instance
            return path if path else None
        elif platform.system() == "Linux":
            try:
                result = subprocess.run(
                    ['zenity', '--file-selection', '--directory', '--title', title],
                    capture_output=True,
                    text=True
                )
                return result.stdout.strip() if result.returncode == 0 else None
            except FileNotFoundError:
                st.error("Zenity is not installed. Please install it for file dialog support.")
                return None
        elif platform.system() == "Darwin":  # macOS
            try:
                result = subprocess.run(
                    ['osascript', '-e', f'choose folder with prompt "{title}"'],
                    capture_output=True,
                    text=True
                )
                return result.stdout.strip() if result.returncode == 0 else None
            except FileNotFoundError:
                st.error("AppleScript is not available for file dialog support.")
                return None
    except Exception as e:
        st.error(f"Error opening file dialog: {str(e)}")
        return None
    return None

def create_zip_file(files_dict):
    """Create a zip file containing all processed files"""
    zip_buffer = BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
        for filename, content in files_dict.items():
            zip_file.writestr(filename, content)
    return zip_buffer

# 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
st.header(f"{file_type} File Processing")

col1, col2 = st.columns(2)

with col1:
    # Input method selection
    input_method = st.radio(
        "Choose input method:",
        ("Upload Files", "Select Directory"),
        help="Upload files directly or process from local directory"
    )
    
    if input_method == "Upload Files":
        uploaded_files = st.file_uploader(
            f"Upload {file_type} files", 
            type=["xlsx", "xlsm"] if file_type == "Excel" else ["docx", "docm"],
            accept_multiple_files=True,
            help=f"You can upload multiple {file_type} files"
        )
    else:  # Select Directory
        # Source Directory
        source_dir = st.text_input("Source Directory Path", 
                          value=st.session_state.get('source_dir', ''),
                          key='source_dir_input',
                          help="Enter the full path to the directory containing files to process")

        col_src1, col_src2 = st.columns([1, 4])
        with col_src1:
            browse_clicked = st.button("Browse", key='source_browse')
            if browse_clicked:
                try:
                    if platform.system() == "Darwin":  # macOS
                        try:
                            result = subprocess.run(
                                ['osascript', '-e', 'choose folder with prompt "Select Source Directory"'],
                                capture_output=True,
                                text=True
                            )
                            if result.returncode == 0:
                                path = result.stdout.strip()
                                st.session_state['source_dir_selected'] = path
                                st.rerun()
                        except Exception as e:
                            st.error(f"Error using AppleScript: {str(e)}")
                    else:  # Windows or Linux
                        import tkinter as tk
                        from tkinter import filedialog
                        import threading
                        import queue

                        # Create a queue to pass the selected path between threads
                        path_queue = queue.Queue()

                        def create_file_dialog():
                            root = tk.Tk()
                            root.withdraw()
                            root.wm_attributes('-topmost', 1)
                            path = filedialog.askdirectory(title="Select Source Directory")
                            path_queue.put(path)
                            root.destroy()

                        # Create and start the dialog in a new thread
                        dialog_thread = threading.Thread(target=create_file_dialog)
                        dialog_thread.start()
                        dialog_thread.join()  # Wait for the thread to complete

                        # Get the selected path from the queue
                        selected_path = path_queue.get()

                        if selected_path:
                            st.session_state['source_dir_selected'] = selected_path
                            st.rerun()
                except Exception as e:
                    st.error(f"Error selecting directory: {str(e)}")

        # Initialize the source directory input with the selected path
        if 'source_dir_selected' in st.session_state:
            source_dir = st.session_state['source_dir_selected']

        # Add a check and display for source directory status
        if source_dir:
            if os.path.exists(source_dir):
                st.success(f"✅ Source directory exists: {source_dir}")
            else:
                st.error(f"❌ Source directory not found: {source_dir}")

        # Destination Directory        
        dest_dir = st.text_input("Destination Directory Path",
                        value=st.session_state.get('dest_dir', ''),
                        key='dest_dir_input',
                        help="Enter the full path where processed files will be saved")

        col_dest1, col_dest2 = st.columns([1, 4])
        with col_dest1:
            browse_clicked = st.button("Browse", key='dest_browse')
            if browse_clicked:
                try:
                    if platform.system() == "Darwin":  # macOS
                        try:
                            result = subprocess.run(
                                ['osascript', '-e', 'choose folder with prompt "Select Destination Directory"'],
                                capture_output=True,
                                text=True
                            )
                            if result.returncode == 0:
                                path = result.stdout.strip()
                                st.session_state['dest_dir_selected'] = path
                                st.rerun()
                        except Exception as e:
                            st.error(f"Error using AppleScript: {str(e)}")
                    else:  # Windows or Linux
                        import tkinter as tk
                        from tkinter import filedialog
                        import threading
                        import queue

                        # Create a queue to pass the selected path between threads
                        path_queue = queue.Queue()

                        def create_file_dialog():
                            root = tk.Tk()
                            root.withdraw()
                            root.wm_attributes('-topmost', 1)
                            path = filedialog.askdirectory(title="Select Destination Directory")
                            path_queue.put(path)
                            root.destroy()

                        # Create and start the dialog in a new thread
                        dialog_thread = threading.Thread(target=create_file_dialog)
                        dialog_thread.start()
                        dialog_thread.join()  # Wait for the thread to complete

                        # Get the selected path from the queue
                        selected_path = path_queue.get()

                        if selected_path:
                            st.session_state['dest_dir_selected'] = selected_path
                            st.rerun()
                except Exception as e:
                    st.error(f"Error selecting directory: {str(e)}")

        # Initialize the destination directory input with the selected path
        if 'dest_dir_selected' in st.session_state:
            dest_dir = st.session_state['dest_dir_selected']

        # Add a check and display for destination directory status
        if dest_dir:
            dest_parent = os.path.dirname(dest_dir)
            if os.path.exists(dest_parent):
                if os.path.exists(dest_dir):
                    st.success(f"✅ Destination directory exists: {dest_dir}")
                else:
                    st.info(f"i️ Destination directory will be created: {dest_dir}")
            else:
                st.error(f"❌ Parent directory not found: {dest_parent}")

with col2:
    if file_type == "Excel":
        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":
            if input_method == "Upload Files":
                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:  # Select Directory
                password_path = st.text_input("Password File Path", 
                                            help="Enter the full path to the text file containing passwords",
                                            value=st.session_state.get('password_path', ''))

                # Use direct file uploader for password file
                password_file_upload = st.file_uploader("Or upload password file", 
                                                      type=["txt"],
                                                      key='password_file_selector',
                                                      help="Upload a text file containing passwords")
                if password_file_upload:
                    content = password_file_upload.getvalue().decode()
                    passwords = [line.strip() for line in content.splitlines() if line.strip()]
                    st.info(f"Loaded {len(passwords)} passwords")
                elif password_path and os.path.exists(password_path):
                    with open(password_path, 'r', encoding='utf-8') as pf:
                        passwords = [line.strip() for line in pf if line.strip()]
                    st.info(f"Loaded {len(passwords)} passwords from file")

# Process button and logic
if input_method == "Upload Files" and uploaded_files and st.button("Process Files", type="primary"):
    progress_bar = st.progress(0)
    status_text = st.empty()
    processed_files = {}
    
    for idx, uploaded_file in enumerate(uploaded_files):
        try:
            with st.expander(f"Processing {uploaded_file.name}", expanded=True):
                st.write(f"📝 Processing {uploaded_file.name}...")
                temp_input_path = save_uploaded_file(uploaded_file)
                temp_output_path = f"{temp_input_path}_processed"
                
                if file_type == "Excel":
                    # Pass original filename to copy_excel_file
                    chartsheet_warning = copy_excel_file(temp_input_path, temp_output_path, passwords, original_filename=uploaded_file.name)
                    if chartsheet_warning:
                        warning_msg = f"⚠️ {chartsheet_warning} Please check the processed file against the original to ensure data was copied correctly."
                        st.warning(warning_msg)
                        st.session_state.error_log[uploaded_file.name] = warning_msg
                else:  # Word
                    remove_all_protection_tags(temp_input_path, temp_output_path)
                
                with open(temp_output_path, "rb") as f:
                    processed_file = f.read()
                    processed_files[f"processed_{uploaded_file.name}"] = processed_file
                
                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
                )
                
                os.unlink(temp_input_path)
                os.unlink(temp_output_path)
                st.success("✅ Processing complete!")
            
            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:
            error_msg = f"❌ Error processing {uploaded_file.name}: {str(e)}"
            st.error(error_msg)
            st.session_state.error_log[uploaded_file.name] = str(e)
    
    if len(processed_files) > 1:
        zip_buffer = create_zip_file(processed_files)
        st.download_button(
            label="⬇️ Download all files as ZIP",
            data=zip_buffer.getvalue(),
            file_name="processed_files.zip",
            mime="application/zip",
        )

elif input_method == "Select Directory" and source_dir and dest_dir and st.button("Process Directory", type="primary"):
    if not os.path.exists(source_dir):
        st.error(f"Source directory does not exist: {source_dir}")
    elif not os.path.exists(os.path.dirname(dest_dir)):
        st.error(f"Parent of destination directory does not exist: {os.path.dirname(dest_dir)}")
    else:
        os.makedirs(dest_dir, exist_ok=True)
        
        # Get all files recursively
        all_files = []
        files_to_process = []
        
        for root, _, files in os.walk(source_dir):
            for file in files:
                full_path = os.path.join(root, file)
                file_lower = file.lower()
                
                if file_type == "Excel" and file_lower.endswith(('.xlsx', '.xlsm')):
                    files_to_process.append(full_path)
                elif file_type == "Word" and file_lower.endswith(('.docx', '.docm')):
                    files_to_process.append(full_path)
                else:
                    all_files.append(full_path)
        
        if not files_to_process:
            st.warning(f"No {file_type} files found in the source directory")
        
        total_files = len(files_to_process) + len(all_files)
        progress_bar = st.progress(0)
        status_text = st.empty()
        files_processed = 0
        
        # Process Office files
        for source_path in files_to_process:
            try:
                relative_path = os.path.relpath(source_path, source_dir)
                dest_path = os.path.join(dest_dir, relative_path)
                
                with st.expander(f"Processing {relative_path}", expanded=True):
                    st.write(f"📝 Processing {relative_path}...")
                    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
                    
                    if file_type == "Excel":
                        chartsheet_warning = copy_excel_file(source_path, dest_path, passwords)
                        if chartsheet_warning:
                            warning_msg = f"⚠️ {chartsheet_warning} Please check the processed file against the original to ensure data was copied correctly."
                            st.warning(warning_msg)
                            st.session_state.error_log[relative_path] = warning_msg
                    else:  # Word
                        remove_all_protection_tags(source_path, dest_path)
                    
                    st.success("✅ Processing complete!")
                
                files_processed += 1
                progress = files_processed / total_files
                progress_bar.progress(progress)
                status_text.text(f"Processed {files_processed} of {total_files} files")
                
            except Exception as e:
                error_msg = f"❌ Error processing {relative_path}: {str(e)}"
                st.error(error_msg)
                st.session_state.error_log[source_path] = str(e)
        
        # Copy all other files
        with st.expander("Copying other files", expanded=True):
            for source_path in all_files:
                try:
                    relative_path = os.path.relpath(source_path, source_dir)
                    dest_path = os.path.join(dest_dir, relative_path)
                    
                    # Create destination directory if it doesn't exist
                    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
                    
                    # Copy the file
                    import shutil
                    shutil.copy2(source_path, dest_path)
                    
                    files_processed += 1
                    progress = files_processed / total_files
                    progress_bar.progress(progress)
                    status_text.text(f"Processed {files_processed} of {total_files} files")
                    
                except Exception as e:
                    error_msg = f"❌ Error copying {relative_path}: {str(e)}"
                    st.error(error_msg)
                    st.session_state.error_log[source_path] = str(e)
        
        st.success(f"✨ All files processed and saved to: {dest_dir}")
        if len(all_files) > 0:
            st.info(f"📁 Copied {len(all_files)} additional files to maintain folder structure")

# Error Log Section
st.markdown("---")
st.header("Error Log")

if 'error_log' not in st.session_state:
    st.session_state.error_log = {}

# Display error log if there are errors
if st.session_state.error_log:
    st.error("The following errors were encountered:")
    error_text = "\n\n".join([f"File: {path}\nError: {error}" for path, error in st.session_state.error_log.items()])
    st.text_area("Error Details", error_text, height=200)
    
    # Add copy button with JavaScript to copy to clipboard
    copy_button = st.button("Copy Error Log")
    if copy_button:
        # Create a JavaScript function to copy the text
        # Escape any backticks in the error text first
        escaped_error_text = error_text.replace('`', '\u0060')
        
        js_code = f"""
            <script>
                async function copyToClipboard() {{
                    try {{
                        const text = `{escaped_error_text}`;
                        await navigator.clipboard.writeText(text);
                        return true;
                    }} catch (err) {{
                        console.error('Failed to copy text: ', err);
                        return false;
                    }}
                }}
                copyToClipboard();
            </script>
        """
        st.components.v1.html(js_code, height=0)
        st.success("✅ Error log copied to clipboard!")
    
    # Add clear button
    if st.button("Clear Error Log"):
        st.session_state.error_log = {}
        st.rerun()
else:
    st.success("No errors encountered in current session")

# Footer
st.sidebar.markdown("---")
st.sidebar.markdown("### Instructions")
st.sidebar.markdown("""
1. Select file type (Excel or Word)
2. Choose input method:
   - Upload Files: Process files via web upload
   - Select Directory: Process files from local directories
3. For Excel files, set password options if needed
4. Click Process button
5. Monitor progress and download processed files
""")

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)

Upload Limits:
- Individual files: up to 200MB each
- Total upload size: up to 800MB per session

For larger files or bulk processing, use the 'Select Directory' option to process files locally.

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