#!/usr/bin/env python3
"""
Web-based GUI Anonymizer - Browser interface for image anonymization pipeline

Features:
- Folder browser to select images
- Side-by-side before/after comparison
- Live chat log showing agent conversations
- Process queue for batch processing
- Works on remote servers (no X11 needed)
- Always uses full multi-agent pipeline (all PII types)

Usage:
    python gui_web_anonymizer.py
    
    Then open the URL shown in your browser (e.g., http://localhost:7860)
    
Note:
    GUI always uses the full pipeline for comprehensive anonymization.
    For faster persons-only batch processing, use batch_process.py instead.
    
Dependencies:
    pip install gradio pillow
"""

import gradio as gr
from PIL import Image
import json
import subprocess
from pathlib import Path
import threading
from queue import Queue, Empty
import time
from datetime import datetime
import sys
import os
from typing import List, Tuple, Optional
import io

class WebAnonymizerGUI:
    def __init__(self):
        self.current_folder = None
        self.image_files = []
        self.processing = False
        self.chat_history = []
        self.dataset_paths = {
            "Market-1501": "data/Market-1501-v15.09.15/bounding_box_test",
            "Detected (bounding_box_test)": "data/detected/bounding_box_test",
            "Detected (bounding_box_train)": "data/detected/bounding_box_train",
            "Detected (query)": "data/detected/query",
            "CityScapes (train)": "data/CityScapes/leftImg8bit/train",
            "CityScapes (val)": "data/CityScapes/leftImg8bit/val",
            "CityScapes (test)": "data/CityScapes/leftImg8bit/test",
            "Mapillary (training)": "data/Mapillary-Vistas-v2.0/training/images",
            "Mapillary (validation)": "data/Mapillary-Vistas-v2.0/validation/images",
            "test_data": "test_data/other",
        }
        
    def list_datasets(self):
        """List available datasets with their status"""
        datasets = []
        for name, path in self.dataset_paths.items():
            path_obj = Path(path)
            if not path_obj.exists():
                continue
            
            # Count images (recursively for nested structures)
            img_count = len(list(path_obj.rglob('*.[jpJP][npNP]*[gG]')))  # jpg, jpeg, png
            
            # Count processed images in artifacts
            artifacts_path = Path("artifacts") / Path(path)
            processed_count = len(list(artifacts_path.rglob('*.jpeg'))) if artifacts_path.exists() else 0
            
            if img_count > 0:
                status = f"{processed_count}/{img_count} processed" if processed_count > 0 else "not started"
                datasets.append(f"{name} ({status})")
        
        return sorted(datasets)
    
    def list_folders(self, base_path: str = "test_data"):
        """List available folders containing images (legacy method for test_data)"""
        base = Path(base_path)
        if not base.exists():
            return []
        
        folders = []
        for item in base.iterdir():
            if item.is_dir():
                # Count images in folder
                img_count = len([f for f in item.iterdir() 
                               if f.suffix.lower() in ('.png', '.jpg', '.jpeg', '.bmp', '.tiff')])
                if img_count > 0:
                    folders.append(f"{item.name} ({img_count} images)")
        
        return sorted(folders)
    
    def load_folder_images(self, dataset_info: str):
        """Load images from selected dataset"""
        if not dataset_info:
            return gr.Radio(choices=[], value=None), "Please select a dataset"
        
        # Extract dataset name - need to handle both "Dataset (status)" and "Dataset (subfolder) (status)" formats
        # For "Detected (bounding_box_test) (22/5332 processed)", we want "Detected (bounding_box_test)"
        # Split by ' (' and rejoin all but the last part (which is the status)
        parts = dataset_info.split(' (')
        if len(parts) > 1:
            # Remove the last part (status) and reconstruct the name
            dataset_name = ' ('.join(parts[:-1])
        else:
            dataset_name = dataset_info
        
        # Find matching path
        folder_path = None
        for name, path in self.dataset_paths.items():
            if name == dataset_name:
                folder_path = Path(path)
                break
        
        if not folder_path or not folder_path.exists():
            return gr.Radio(choices=[], value=None), f"Dataset path not found: {dataset_name}"
        
        self.current_folder = folder_path
        
        # Find all image files (recursively for nested structures like CityScapes)
        extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.tiff')
        self.image_files = []
        
        # For CityScapes and Mapillary, search recursively
        if "CityScapes" in dataset_name or "Mapillary" in dataset_name:
            self.image_files = sorted([
                f for f in folder_path.rglob('*')
                if f.is_file() and f.suffix.lower() in extensions
            ])
        else:
            # For Market-1501 and test_data, just check the folder
            self.image_files = sorted([
                f for f in folder_path.iterdir()
                if f.is_file() and f.suffix.lower() in extensions
            ])
        
        # Create file list with status
        file_list = []
        for img_file in self.image_files[:500]:  # Limit to first 500 for UI performance
            # Get relative path for display
            try:
                rel_path = img_file.relative_to(folder_path)
                display_name = str(rel_path)
            except ValueError:
                display_name = img_file.name
            
            # Check if output exists in artifacts
            # Convert img_file to absolute path, then get relative to cwd
            try:
                img_abs = img_file.resolve()
                cwd_abs = Path.cwd().resolve()
                rel_to_cwd = img_abs.relative_to(cwd_abs)
                artifacts_base = Path("artifacts") / rel_to_cwd
                output_path = artifacts_base.parent / f"{img_file.stem}.jpeg"
            except ValueError:
                # Fallback: just use the image's relative path from folder
                output_path = Path("artifacts") / folder_path / f"{img_file.stem}.jpeg"
                
            status = "✅" if output_path.exists() else "⏳"
            file_list.append(f"{status} {display_name}")
        
        total_msg = f"Loaded {len(self.image_files)} images from {dataset_name}"
        if len(self.image_files) > 500:
            total_msg += f" (showing first 500)"
        
        # Return update for Radio component
        return gr.Radio(choices=file_list, value=None), total_msg
    
    def get_image_pair(self, selected_file: str) -> Tuple[Optional[Image.Image], Optional[Image.Image], str]:
        """Get before/after images for selected file"""
        if not selected_file or not self.image_files:
            return None, None, "No file selected"
        
        # Extract filename (remove status emoji)
        display_name = selected_file.split(' ', 1)[1] if ' ' in selected_file else selected_file
        
        # Find file by matching display name
        img_path = None
        for f in self.image_files:
            # Check both full relative path and just the name
            try:
                rel_path = f.relative_to(self.current_folder)
                if str(rel_path) == display_name or f.name == display_name:
                    img_path = f
                    break
            except ValueError:
                if f.name == display_name:
                    img_path = f
                    break
        
        if not img_path:
            return None, None, f"File not found: {display_name}"
        
        # Load before image
        try:
            before_img = Image.open(img_path)
        except Exception as e:
            return None, None, f"Error loading image: {e}"
        
        # Determine output path in artifacts (mirrors source structure)
        try:
            img_abs = img_path.resolve()
            cwd_abs = Path.cwd().resolve()
            rel_to_cwd = img_abs.relative_to(cwd_abs)
            artifacts_base = Path("artifacts") / rel_to_cwd
            output_path = artifacts_base.parent / f"{img_path.stem}.jpeg"
        except ValueError:
            # Fallback if relative path fails
            output_path = Path("artifacts") / self.current_folder.name / f"{img_path.stem}.jpeg"
        
        after_img = None
        info = f"**File:** {display_name}\n"
        info += f"**Size:** {before_img.size[0]}x{before_img.size[1]} px\n"
        info += f"**Path:** {img_path}\n"
        
        if output_path.exists():
            try:
                after_img = Image.open(output_path)
                info += f"\n**Status:** ✅ Processed\n"
                info += f"**Output:** {output_path}\n"
                
                # Load metadata (JSON file has same name as output)
                json_path = output_path.with_suffix('.json')
                if json_path.exists():
                    with open(json_path) as f:
                        meta = json.load(f)
                    info += f"**PII Instances:** {len(meta.get('instances', []))}\n"
                    info += f"**Operations:** {len(meta.get('steps', []))}\n"
                    timing = meta.get('timing', {})
                    if 'workflow_seconds' in timing:
                        info += f"**Processing Time:** {timing['workflow_seconds']:.1f}s\n"
                    audit = meta.get('audit', {})
                    if audit.get('ok'):
                        info += "**Audit:** ✅ Passed\n"
                    else:
                        info += f"**Audit:** ⚠️ {audit.get('error', 'Unknown issue')}\n"
            except Exception as e:
                info += f"\n**Error loading result:** {e}\n"
        else:
            info += f"\n**Status:** ⏳ Not processed\n"
            info += f"**Expected output:** {output_path}\n"
        
        return before_img, after_img, info
    
    def process_single_image(self, selected_file: str):
        """Process a single image with real-time chat streaming"""
        if not selected_file or not self.image_files:
            yield None, "❌ No file selected"
            return
        
        if self.processing:
            yield None, "⏳ Already processing..."
            return
        
        # Extract display name (remove status emoji)
        display_name = selected_file.split(' ', 1)[1] if ' ' in selected_file else selected_file
        
        # Find file by matching display name
        img_path = None
        for f in self.image_files:
            try:
                rel_path = f.relative_to(self.current_folder)
                if str(rel_path) == display_name or f.name == display_name:
                    img_path = f
                    break
            except ValueError:
                if f.name == display_name:
                    img_path = f
                    break
        
        if not img_path:
            yield None, f"❌ File not found: {display_name}"
            return
        
        self.processing = True
        self.chat_history = []
        
        yield None, f"🚀 Starting processing: {display_name}\n\n"
        
        # Run full pipeline for all images in GUI
        try:
            cmd = [
                sys.executable,
                "-m", "mcp.clients.agents.main",
                str(img_path)
            ]
            
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,  # Merge stderr into stdout
                text=True,
                bufsize=1,
                cwd=str(Path.cwd())
            )
            
            # Read output line by line and stream to UI (captures both stdout and stderr)
            output_lines = []
            for i, line in enumerate(process.stdout):
                line = line.rstrip()  # Keep structure but remove trailing whitespace
                if line:
                    output_lines.append(line)
                    formatted_line = self.format_chat_line(line)
                    self.chat_history.append(formatted_line)
                    
                    # Yield updated chat after each line (real-time streaming)
                    chat_text = "\n".join(self.chat_history)
                    yield None, chat_text
            
            process.wait()
            
            self.processing = False
            
            # Load result
            try:
                img_abs = img_path.resolve()
                cwd_abs = Path.cwd().resolve()
                rel_to_cwd = img_abs.relative_to(cwd_abs)
                artifacts_base = Path("artifacts") / rel_to_cwd
                output_path = artifacts_base.parent / f"{img_path.stem}.jpeg"
            except ValueError:
                output_path = Path("artifacts") / self.current_folder.name / f"{img_path.stem}.jpeg"
            
            if process.returncode == 0 and output_path.exists():
                result_img = Image.open(output_path)
                chat_text = "\n".join(self.chat_history)
                final_message = f"✅ Successfully processed {display_name}\n\n{chat_text}"
                yield result_img, final_message
            else:
                chat_text = "\n".join(self.chat_history)
                final_message = f"❌ Processing failed\n\n{chat_text}"
                yield None, final_message
                
        except Exception as e:
            self.processing = False
            yield None, f"❌ Error: {str(e)}"
    
    def format_chat_line(self, line: str) -> str:
        """Format a log line with emojis for readability"""
        # Agent conversations
        if "UserProxyAgent (to chat_manager):" in line:
            return f"👤 {line}"
        elif "OrchestratorAgent (to chat_manager):" in line:
            return f"🎯 {line}"
        elif "AuditorAgent (to chat_manager):" in line:
            return f"🔍 {line}"
        elif "DetectionAgent (to chat_manager):" in line:
            return f"👁️ {line}"
        elif "GenerativeAgent (to chat_manager):" in line:
            return f"🎨 {line}"
        elif "TechnicianAgent (to chat_manager):" in line:
            return f"🔧 {line}"
        # Tool calls
        elif "***** Suggested tool call" in line:
            return f"🔧 {line}"
        elif "***** Response from calling tool" in line:
            return f"📦 {line}"
        elif "EXECUTING FUNCTION" in line:
            return f"⚙️ {line}"
        elif "EXECUTED FUNCTION" in line:
            return f"✅ {line}"
        # Status messages
        elif "Next speaker:" in line:
            return f"➡️ {line}"
        elif "USING AUTO REPLY" in line:
            return f"🤖 {line}"
        elif "TERMINATING RUN" in line:
            return f"⏸️ {line}"
        elif "PIPELINE COMPLETE" in line:
            return f"🎉 {line}"
        # Errors and warnings
        elif "ERROR" in line or "Failed" in line:
            return f"❌ {line}"
        elif "WARNING" in line:
            return f"⚠️ {line}"
        # Success
        elif "Successfully" in line or "complete:" in line.lower():
            return f"✅ {line}"
        # Processing logs
        elif "[Generative]" in line or "[Detector]" in line or "[Inpainter]" in line:
            return f"📝 {line}"
        else:
            return line
    
    def process_all_images(self, dataset_info: str):
        """Process all unprocessed images in dataset - yields real-time updates"""
        if not dataset_info:
            yield "❌ No dataset selected"
            return
        
        if self.processing:
            yield "⏳ Already processing..."
            return
        
        if not self.image_files:
            yield "❌ No images loaded"
            return
        
        # Find unprocessed images
        # Extract dataset name - handle "Dataset (subfolder) (status)" format
        parts = dataset_info.split(' (')
        if len(parts) > 1:
            dataset_name = ' ('.join(parts[:-1])
        else:
            dataset_name = dataset_info
        unprocessed = []
        
        for img_file in self.image_files:
            try:
                img_abs = img_file.resolve()
                cwd_abs = Path.cwd().resolve()
                rel_to_cwd = img_abs.relative_to(cwd_abs)
                artifacts_base = Path("artifacts") / rel_to_cwd
                output_path = artifacts_base.parent / f"{img_file.stem}.jpeg"
            except ValueError:
                output_path = Path("artifacts") / self.current_folder.name / f"{img_file.stem}.jpeg"
                
            if not output_path.exists():
                unprocessed.append(img_file)
        
        if not unprocessed:
            yield "✅ All images already processed!"
            return
        
        total = len(unprocessed)
        results = []
        
        for i, img_file in enumerate(unprocessed):
            self.processing = True
            
            # Get display name
            try:
                display_name = str(img_file.relative_to(self.current_folder))
            except ValueError:
                display_name = img_file.name
            
            # Clear chat for new image
            self.chat_history = []
            status_msg = f"🎬 Processing image {i+1}/{total}: {display_name}\n\n"
            yield status_msg + "\n".join(results)
            
            try:
                # Always use full pipeline in GUI
                cmd = [
                    sys.executable,
                    "-m", "mcp.clients.agents.main",
                    str(img_file)
                ]
                
                process = subprocess.Popen(
                    cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,  # Merge stderr into stdout
                    text=True,
                    bufsize=1,
                    cwd=str(Path.cwd())
                )
                
                # Stream chat output while processing (captures both stdout and stderr)
                for line in process.stdout:
                    line = line.rstrip()  # Keep structure but remove trailing whitespace
                    if line:
                        formatted_line = self.format_chat_line(line)
                        self.chat_history.append(formatted_line)
                        chat_text = "\n".join(self.chat_history)
                        
                        # Yield status + previous results + current chat
                        current_status = f"🎬 Processing image {i+1}/{total}: {display_name}\n\n"
                        if results:
                            current_status += "\n".join(results) + "\n\n"
                        current_status += "--- Current Agent Conversation ---\n" + chat_text
                        
                        yield current_status
                
                process.wait()
                
                try:
                    img_abs = img_file.resolve()
                    cwd_abs = Path.cwd().resolve()
                    rel_to_cwd = img_abs.relative_to(cwd_abs)
                    artifacts_base = Path("artifacts") / rel_to_cwd
                    output_path = artifacts_base.parent / f"{img_file.stem}.jpeg"
                except ValueError:
                    output_path = Path("artifacts") / self.current_folder.name / f"{img_file.stem}.jpeg"
                    
                if process.returncode == 0 and output_path.exists():
                    results.append(f"✅ {display_name}")
                else:
                    results.append(f"❌ {display_name}")
                    
            except Exception as e:
                results.append(f"❌ {display_name}: {str(e)}")
            
            self.processing = False
        
        summary = f"Batch processing complete: {total} images\n\n" + "\n".join(results)
        yield summary


def create_interface():
    """Create the Gradio interface"""
    
    app = WebAnonymizerGUI()
    
    with gr.Blocks(title="Image Anonymization Pipeline", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🖼️ Image Anonymization Pipeline
        ### Interactive browser-based interface with multi-agent AI
        
        Browse and process images from Market-1501, Detected, CityScapes, Mapillary, and test datasets.
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📁 Step 1: Select Dataset")
                dataset_dropdown = gr.Dropdown(
                    choices=app.list_datasets(),
                    label="Available Datasets",
                    interactive=True
                )
                refresh_btn = gr.Button("🔄 Refresh Datasets", size="sm")
                
                gr.Markdown("### 📄 Step 2: Select Image")
                file_list = gr.Radio(
                    choices=[],
                    label="Images (✅=processed, ⏳=pending)",
                    interactive=True
                )
                
                load_status = gr.Textbox(label="Status", lines=2, interactive=False)
                
            with gr.Column(scale=2):
                gr.Markdown("### 🖼️ Before / After Comparison")
                
                with gr.Row():
                    before_img = gr.Image(label="Before", type="pil", interactive=False)
                    after_img = gr.Image(label="After", type="pil", interactive=False)
                
                image_info = gr.Markdown("Select an image to view details")
                
                with gr.Row():
                    process_btn = gr.Button("▶️ Process Selected Image", variant="primary", size="lg")
                    batch_btn = gr.Button("⏩ Process All Unprocessed", variant="secondary", size="lg")
        
        gr.Markdown("### 💬 Pipeline Output & Processing Log")
        chat_output = gr.Textbox(
            label="Real-time Processing Output",
            lines=20,
            max_lines=30,
            interactive=False,
            placeholder="Processing output will appear here...\n\nYou'll see:\n[Pipeline] Phase 1 - Deterministic preprocessing\n[Detection] Person detection with YOLOv8\n[Generative] Person inpainting with Stable Diffusion\n✅ Success messages\n❌ Error messages"
        )
        
        with gr.Row():
            export_btn = gr.Button("💾 Export Results", size="sm")
            clear_btn = gr.Button("🗑️ Clear Log", size="sm")
        
        # Event handlers
        refresh_btn.click(
            fn=lambda: gr.Dropdown(choices=app.list_datasets()),
            outputs=dataset_dropdown
        )
        
        dataset_dropdown.change(
            fn=app.load_folder_images,
            inputs=dataset_dropdown,
            outputs=[file_list, load_status],
            show_progress=True
        )
        
        file_list.select(
            fn=app.get_image_pair,
            inputs=file_list,
            outputs=[before_img, after_img, image_info]
        )
        
        process_btn.click(
            fn=app.process_single_image,
            inputs=file_list,
            outputs=[after_img, chat_output]
        )
        
        batch_btn.click(
            fn=app.process_all_images,
            inputs=dataset_dropdown,
            outputs=chat_output
        )
        
        clear_btn.click(
            fn=lambda: "",
            outputs=chat_output
        )
        
        gr.Markdown("""
        ---
        ### 📖 Quick Guide
        
        1. **Select a dataset** from the dropdown:
           - **Market-1501**: Person re-identification dataset (bounding_box_test)
           - **Detected**: Detected person crops (bounding_box_test/train, query)
           - **CityScapes**: Urban street scenes (train/val/test splits)
           - **Mapillary**: Street-level imagery (training/validation)
           - **test_data**: Custom test images
        
        2. **Choose an image** from the list (showing up to 500 images)
        3. **View before/after** in the side-by-side comparison
        4. **Process images:**
           - Single: Select an image and click "Process Selected"
           - Batch: Click "Process All Unprocessed" to handle remaining images
        5. **Monitor progress** in the real-time output log
        6. **Export results** when complete
        
        **Status Indicators:**
        - ✅ = Already processed (result available in artifacts/)
        - ⏳ = Pending processing
        
        **Processing Mode:**
        - **GUI always uses full multi-agent pipeline** for all datasets
        - Detects and anonymizes: persons, license plates, traffic signs, and other PII
        - Processing time: ~30-60s per image depending on complexity
        - For faster persons-only batch processing, use command-line `batch_process.py`
        
        **Output Location:**
        - Results saved to `artifacts/{dataset_path}/` mirroring source structure
        - Each output includes: `.jpeg` (anonymized image), `.json` (metadata), `.log` (processing log)
        """)
    
    return demo


if __name__ == "__main__":
    print("🚀 Starting Web-based Anonymizer GUI...")
    print("")
    print("📖 This will open in your browser automatically")
    print("   If not, copy the URL from below and open it manually")
    print("")
    
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",  # Listen on all interfaces (for remote access)
        server_port=7860,
        share=False,  # Set to True to create a public URL
        inbrowser=False  # Don't auto-open browser (may not work on remote server)
    )
