import gradio as gr
import os
import zipfile
import tempfile
import shutil
from pathlib import Path
import mimetypes

# Create uploads directory
UPLOADS_FOLDER = "uploads"
os.makedirs(UPLOADS_FOLDER, exist_ok=True)

def process_file(file_obj, file_name, action):
    """
    Process uploaded file: return original or create zip
    """
    if not file_obj:
        return None, "❌ No file uploaded"
    
    try:
        # Save uploaded file temporarily
        temp_dir = tempfile.mkdtemp()
        original_path = os.path.join(temp_dir, file_name or os.path.basename(file_obj.name))
        
        # Copy file to temp location
        shutil.copy2(file_obj.name, original_path)
        
        if action == "original":
            # Return original file
            return original_path, f"✅ Ready to download: {os.path.basename(original_path)}"
        
        elif action == "zip":
            # Create zip file
            base_name = os.path.splitext(os.path.basename(original_path))[0]
            zip_path = os.path.join(temp_dir, f"{base_name}.zip")
            
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                zipf.write(original_path, os.path.basename(original_path))
            
            return zip_path, f"✅ Zipped as: {os.path.basename(zip_path)}"
        
        elif action == "zip_with_password":
            # Create password-protected zip (requires pyminizip)
            try:
                import pyminizip
                base_name = os.path.splitext(os.path.basename(original_path))[0]
                zip_path = os.path.join(temp_dir, f"{base_name}_protected.zip")
                
                # Compress with password
                pyminizip.compress(
                    original_path, 
                    None, 
                    zip_path, 
                    "password123",  # Default password
                    5  # Compression level
                )
                return zip_path, f"✅ Password-protected zip created (password: password123)"
            
            except ImportError:
                return None, "❌ pyminizip not installed. Install with: pip install pyminizip"
        
        else:
            return None, "❌ Invalid action selected"
    
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

def process_multiple_files(files, action):
    """
    Process multiple uploaded files
    """
    if not files:
        return None, "❌ No files uploaded"
    
    try:
        temp_dir = tempfile.mkdtemp()
        
        if action == "individual":
            # Create zip containing all files
            zip_path = os.path.join(temp_dir, "files.zip")
            
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for file_obj in files:
                    file_path = file_obj.name
                    zipf.write(file_path, os.path.basename(file_path))
            
            return zip_path, f"✅ Created zip with {len(files)} files"
        
        elif action == "separate":
            # Create separate zips for each file
            zip_path = os.path.join(temp_dir, "separate_files.zip")
            
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for file_obj in files:
                    file_path = file_obj.name
                    # Create individual zip first
                    individual_zip = os.path.join(temp_dir, f"{os.path.splitext(os.path.basename(file_path))[0]}.zip")
                    with zipfile.ZipFile(individual_zip, 'w', zipfile.ZIP_DEFLATED) as ind_zip:
                        ind_zip.write(file_path, os.path.basename(file_path))
                    # Add individual zip to main zip
                    zipf.write(individual_zip, os.path.basename(individual_zip))
            
            return zip_path, f"✅ Created zip with {len(files)} individual zips"
        
        else:
            return None, "❌ Invalid action selected"
    
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

def get_file_info(file_obj):
    """Get information about uploaded file"""
    if not file_obj:
        return "No file selected"
    
    try:
        file_path = file_obj.name
        file_size = os.path.getsize(file_path)
        file_ext = os.path.splitext(file_path)[1].lower()
        
        # Get file type
        mime_type, _ = mimetypes.guess_type(file_path)
        file_type = mime_type.split('/')[0] if mime_type else "Unknown"
        
        info = f"""
        📄 **File Information:**
        
        **Name:** {os.path.basename(file_path)}
        **Size:** {file_size:,} bytes ({file_size/1024:.2f} KB)
        **Extension:** {file_ext}
        **Type:** {file_type}
        **Path:** {file_path}
        """
        
        return info
    
    except Exception as e:
        return f"❌ Error getting file info: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="File Upload & Download Manager", theme=gr.themes.Soft()) as iface:
    
    gr.Markdown("# 📁 File Upload & Download Manager")
    gr.Markdown("Upload files and download them as original or zipped")
    
    with gr.Tabs():
        # Single File Tab
        with gr.Tab("Single File"):
            with gr.Row():
                with gr.Column(scale=2):
                    single_file = gr.File(
                        label="Upload a File",
                        file_count="single"
                    )
                    
                    file_name_input = gr.Textbox(
                        label="Custom Filename (optional)",
                        placeholder="Leave empty to keep original name...",
                        info="Enter a custom name for the downloaded file"
                    )
                    
                    single_action = gr.Radio(
                        choices=[
                            ("Download Original", "original"),
                            ("Download as ZIP", "zip"),
                            ("Password-protected ZIP", "zip_with_password")
                        ],
                        label="Select Action",
                        value="original"
                    )
                    
                    single_btn = gr.Button("Process File", variant="primary")
                
                with gr.Column(scale=1):
                    file_info = gr.Markdown(label="File Information")
                    single_status = gr.Textbox(label="Status", interactive=False)
                    single_download = gr.File(label="Download Processed File", interactive=False)
            
            # Update file info when file is uploaded
            single_file.change(
                fn=get_file_info,
                inputs=[single_file],
                outputs=file_info
            )
        
        # Multiple Files Tab
        with gr.Tab("Multiple Files"):
            with gr.Row():
                with gr.Column(scale=2):
                    multi_files = gr.File(
                        label="Upload Multiple Files",
                        file_count="multiple",
                        file_types=["image", "video", "audio", "text", "pdf", ".zip"]
                    )
                    
                    multi_action = gr.Radio(
                        choices=[
                            ("Combine all files into one ZIP", "individual"),
                            ("Create separate ZIPs for each file", "separate")
                        ],
                        label="Select Action",
                        value="individual"
                    )
                    
                    multi_btn = gr.Button("Process Files", variant="primary")
                
                with gr.Column(scale=1):
                    multi_status = gr.Textbox(label="Status", interactive=False)
                    multi_download = gr.File(label="Download Processed Files", interactive=False)
        
        # Batch Processing Tab
        with gr.Tab("Batch Processing"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Upload Multiple Files for Batch Processing")
                    
                    batch_files = gr.File(
                        label="Upload Files",
                        file_count="multiple",
                        file_types=None  # All file types
                    )
                    
                    batch_options = gr.CheckboxGroup(
                        choices=[
                            "Create individual ZIPs",
                            "Create combined ZIP",
                            "Rename with timestamp",
                            "Add to existing ZIP"
                        ],
                        label="Processing Options",
                        value=["Create combined ZIP"]
                    )
                    
                    with gr.Row():
                        batch_format = gr.Dropdown(
                            choices=[".zip", ".7z", ".tar.gz"],
                            value=".zip",
                            label="Archive Format"
                        )
                        
                        compression_level = gr.Slider(
                            minimum=1,
                            maximum=9,
                            value=6,
                            step=1,
                            label="Compression Level"
                        )
                    
                    batch_btn = gr.Button("Process Batch", variant="primary", size="lg")
                
                with gr.Column():
                    batch_status = gr.Textbox(label="Status", interactive=False, lines=3)
                    batch_download = gr.File(label="Download Results", interactive=False)
    
    # Instructions
    with gr.Accordion("📖 Instructions & Features", open=False):
        gr.Markdown("""
        ## How to Use:
        
        ### Single File Tab:
        1. **Upload** a single file
        2. (Optional) Enter a custom filename
        3. **Choose action**: Download original, as ZIP, or password-protected ZIP
        4. Click **Process File**
        
        ### Multiple Files Tab:
        1. **Upload** multiple files (Ctrl+Click to select multiple)
        2. **Choose action**: Combine into one ZIP or create separate ZIPs
        3. Click **Process Files**
        
        ### Batch Processing Tab:
        1. **Upload** multiple files
        2. **Select processing options**
        3. Choose archive format and compression level
        4. Click **Process Batch**
        
        ## Features:
        - ✅ Single file upload and download
        - ✅ Multiple file upload and batch processing
        - ✅ ZIP file creation with compression
        - ✅ Password-protected ZIPs (requires pyminizip)
        - ✅ File information display
        - ✅ Custom filename support
        - ✅ Multiple archive formats
        
        ## Notes:
        - Files are processed in temporary storage
        - Original files are not modified
        - Large files may take time to process
        - Password for protected ZIPs: `password123`
        """)
    
    # Connect events
    single_btn.click(
        fn=process_file,
        inputs=[single_file, file_name_input, single_action],
        outputs=[single_download, single_status]
    )
    
    multi_btn.click(
        fn=process_multiple_files,
        inputs=[multi_files, multi_action],
        outputs=[multi_download, multi_status]
    )
    
    # Batch processing function
    def process_batch(files, options, format_type, compression):
        if not files:
            return None, "❌ No files uploaded"
        
        try:
            temp_dir = tempfile.mkdtemp()
            results = []
            
            # Process based on options
            if "Create combined ZIP" in options:
                zip_name = f"combined{format_type}"
                zip_path = os.path.join(temp_dir, zip_name)
                
                if format_type == ".zip":
                    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED, compresslevel=compression) as zipf:
                        for file_obj in files:
                            file_path = file_obj.name
                            arcname = os.path.basename(file_path)
                            if "Rename with timestamp" in options:
                                import time
                                name, ext = os.path.splitext(arcname)
                                arcname = f"{name}_{int(time.time())}{ext}"
                            zipf.write(file_path, arcname)
                    results.append(zip_path)
            
            if "Create individual ZIPs" in options:
                for file_obj in files:
                    file_path = file_obj.name
                    base_name = os.path.splitext(os.path.basename(file_path))[0]
                    if "Rename with timestamp" in options:
                        import time
                        base_name = f"{base_name}_{int(time.time())}"
                    
                    individual_zip = os.path.join(temp_dir, f"{base_name}{format_type}")
                    
                    if format_type == ".zip":
                        with zipfile.ZipFile(individual_zip, 'w', zipfile.ZIP_DEFLATED, compresslevel=compression) as zipf:
                            zipf.write(file_path, os.path.basename(file_path))
                    
                    results.append(individual_zip)
            
            # If multiple results, create a final zip
            if len(results) > 1:
                final_zip = os.path.join(temp_dir, f"batch_results{format_type}")
                with zipfile.ZipFile(final_zip, 'w', zipfile.ZIP_DEFLATED) as zipf:
                    for result in results:
                        zipf.write(result, os.path.basename(result))
                output_file = final_zip
            elif results:
                output_file = results[0]
            else:
                return None, "⚠️ No processing options selected"
            
            return output_file, f"✅ Processed {len(files)} file(s) with {len(options)} option(s)"
        
        except Exception as e:
            return None, f"❌ Error: {str(e)}"
    
    batch_btn.click(
        fn=process_batch,
        inputs=[batch_files, batch_options, batch_format, compression_level],
        outputs=[batch_download, batch_status]
    )

# Launch the app
if __name__ == "__main__":
    iface.launch(
        server_name="127.0.0.1",
        server_port=7861,  # Different port to avoid conflict
        show_error=True,
        share=False,
        ssr_mode=False
    )