import os
import zipfile
import tempfile
import shutil
import pandas as pd
import numpy as np
import cv2
from PIL import Image
import torch
import torchvision.models as models
import torchvision.transforms as transforms
import urllib.request
import json
import gradio as gr

# Load ImageNet labels
IMAGENET_LABELS = []
try:
    url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
    with urllib.request.urlopen(url, timeout=3) as response:
        IMAGENET_LABELS = [line.decode("utf-8").strip() for line in response.readlines()]
except Exception:
    IMAGENET_LABELS = [f"object_{i}" for i in range(1000)]

# Preprocessing for MobileNet
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

# Global model cache to load lazily
_model = None

def get_model():
    global _model
    if _model is None:
        try:
            # Load lightweight MobileNetV3 Small (approx 10MB)
            _model = models.mobilenet_v3_small(weights=models.MobileNetV3_Small_Weights.DEFAULT)
            _model.eval()
        except Exception as e:
            print(f"Error loading PyTorch model: {e}. Utilizing fallback system.")
            _model = "FALLBACK"
    return _model

# Face detection Cascade
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

def process_single_image(img_path, conf_threshold):
    """Detects faces and predicts object categories in an image."""
    try:
        # Load for OpenCV
        img_cv = cv2.imread(img_path)
        if img_cv is None:
            return 0, ["Error loading image file"], None
            
        h, w, _ = img_cv.shape
        gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
        
        # Face detection
        faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=4, minSize=(30, 30))
        face_count = len(faces)
        
        # Draw bounding boxes
        img_display = img_cv.copy()
        for (x, y, fw, fh) in faces:
            cv2.rectangle(img_display, (x, y), (x+fw, y+fh), (0, 255, 0), max(2, int(w * 0.005)))
            
        # Convert back to RGB for PIL/Gradio
        img_display_rgb = cv2.cvtColor(img_display, cv2.COLOR_BGR2RGB)
        
        # Object detection/Classification
        labels = []
        model = get_model()
        
        if model == "FALLBACK":
            # Rule-based fallback tags using image features
            mean_brightness = np.mean(gray)
            std_brightness = np.std(gray)
            if mean_brightness > 180:
                labels.append("bright_lighting")
            elif mean_brightness < 70:
                labels.append("low_key_lighting")
            if std_brightness > 60:
                labels.append("high_contrast")
            labels.append("visual_media")
        else:
            # PyTorch inference
            img_pil = Image.open(img_path).convert("RGB")
            input_tensor = transform(img_pil).unsqueeze(0)
            with torch.no_grad():
                outputs = model(input_tensor)
                probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
                
            # Get top predictions above confidence threshold
            top_prob, top_catid = torch.topk(probabilities, 5)
            for i in range(5):
                prob = top_prob[i].item()
                if prob >= conf_threshold:
                    class_name = IMAGENET_LABELS[top_catid[i].item()]
                    # Replace underscores with spaces for readability
                    clean_name = class_name.replace("_", " ")
                    labels.append(f"{clean_name} ({prob:.1%})")
                    
            if not labels:
                # Add top 1 as fallback
                class_name = IMAGENET_LABELS[top_catid[0].item()]
                labels.append(class_name.replace("_", " "))
                
        return face_count, labels, img_display_rgb
    except Exception as e:
        print(f"Error processing image {img_path}: {e}")
        return 0, ["Error processing"], None

def initialize_batch(files, conf_threshold):
    """Initializes the batch of images from uploaded files or ZIP."""
    if not files:
        return [], 0, "No files uploaded", None, pd.DataFrame(), None
        
    temp_dir = tempfile.mkdtemp(prefix="visual_labeler_")
    image_paths = []
    
    # Check if a single ZIP was uploaded
    if len(files) == 1 and files[0].name.lower().endswith(".zip"):
        try:
            with zipfile.ZipFile(files[0].name, 'r') as zip_ref:
                zip_ref.extractall(temp_dir)
            # Find all images recursively
            for root, _, filenames in os.walk(temp_dir):
                for filename in filenames:
                    if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.webp')):
                        image_paths.append(os.path.join(root, filename))
        except Exception as e:
            return [], 0, f"Error extracting ZIP: {e}", None, pd.DataFrame(), None
    else:
        # Multiple files uploaded directly
        for f in files:
            dest = os.path.join(temp_dir, os.path.basename(f.name))
            shutil.copy(f.name, dest)
            image_paths.append(dest)
            
    if not image_paths:
        return [], 0, "No valid image files found.", None, pd.DataFrame(), None
        
    # Sort for deterministic order
    image_paths.sort()
    
    # Process first image
    face_count, auto_labels, img_display = process_single_image(image_paths[0], conf_threshold)
    
    # Initialize dataframe
    df_data = []
    for path in image_paths:
        df_data.append({
            "Filename": os.path.basename(path),
            "Auto-detected Tags": "",
            "Face Count": 0,
            "Final Labels (Edited)": ""
        })
    df = pd.DataFrame(df_data)
    
    # Store first result
    df.at[0, "Auto-detected Tags"] = ", ".join(auto_labels)
    df.at[0, "Face Count"] = face_count
    df.at[0, "Final Labels (Edited)"] = ", ".join([l.split(" (")[0] for l in auto_labels])
    
    status_text = f"Successfully loaded {len(image_paths)} images. Displaying image 1 of {len(image_paths)}."
    tags_val = df.at[0, "Final Labels (Edited)"]
    
    return image_paths, 0, status_text, img_display, df, tags_val, temp_dir

def save_and_navigate(direction, current_idx, image_paths, df, tags_val, conf_threshold):
    """Saves current edits and navigates to next/previous image."""
    if not image_paths or current_idx < 0 or current_idx >= len(image_paths):
        return current_idx, "No batch initialized.", None, df, tags_val
        
    # Save edits for the current image
    df.at[current_idx, "Final Labels (Edited)"] = tags_val
    
    # Calculate new index
    new_idx = current_idx + int(direction)
    if new_idx < 0:
        new_idx = 0
    elif new_idx >= len(image_paths):
        new_idx = len(image_paths) - 1
        
    # Load next/prev image
    target_path = image_paths[new_idx]
    
    # Run auto-detection if not done yet
    if not df.at[new_idx, "Auto-detected Tags"]:
        face_count, auto_labels, img_display = process_single_image(target_path, conf_threshold)
        df.at[new_idx, "Auto-detected Tags"] = ", ".join(auto_labels)
        df.at[new_idx, "Face Count"] = face_count
        df.at[new_idx, "Final Labels (Edited)"] = ", ".join([l.split(" (")[0] for l in auto_labels])
    else:
        # Re-read for displaying bounding boxes
        _, _, img_display = process_single_image(target_path, conf_threshold)
        
    status_text = f"Displaying image {new_idx + 1} of {len(image_paths)}."
    next_tags = df.at[new_idx, "Final Labels (Edited)"]
    
    return new_idx, status_text, img_display, df, next_tags

def export_csv(df):
    """Generates a downloadable CSV path of the finalized labeled database."""
    if df.empty:
        return None
    temp_csv = tempfile.NamedTemporaryFile(delete=False, suffix="_labeled_database.csv")
    df.to_csv(temp_csv.name, index=False)
    return temp_csv.name

# Gradient Dark Theme styling
custom_css = """
body { background-color: #0d0f12; color: #e3e6eb; font-family: 'Inter', sans-serif; }
.gradio-container { max-width: 1200px !important; margin: 0 auto !important; }
h1, h2, h3 { color: #ffffff !important; font-weight: 700 !important; }
.pill-tag { background: linear-gradient(135deg, #10b981 0%, #059669 100%) !important; color: white !important; }
.btn-primary { background: linear-gradient(135deg, #10b981 0%, #059669 100%) !important; border: none !important; color: white !important; font-weight: 600 !important; }
.btn-primary:hover { filter: brightness(1.1); }
.btn-secondary { background: #1f2937 !important; border: 1px solid #374151 !important; color: white !important; }
.dataframe-container { background: #111827 !important; border: 1px solid #1f2937 !important; border-radius: 8px; }
"""

with gr.Blocks(theme=gr.themes.Monochrome(), css=custom_css) as demo:
    # State management
    image_paths_state = gr.State([])
    current_idx_state = gr.State(0)
    temp_dir_state = gr.State("")
    
    gr.Markdown(
        """
        # 🕸️ Visual Content Labeler & Batch Database Builder
        ### Expedite visual content analysis, auto-detect faces/objects, review qualitative labels, and build a labeled CSV database in minutes.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=4):
            # Batch Upload Card
            with gr.Card():
                gr.Markdown("### 1. Upload ZIP or Multiple Images")
                files_input = gr.File(
                    file_count="multiple", 
                    label="Upload ZIP archive or drag-and-drop multiple image files",
                    file_types=[".zip", ".png", ".jpg", ".jpeg", ".webp"]
                )
                conf_slider = gr.Slider(
                    minimum=0.1, maximum=1.0, value=0.4, step=0.05,
                    label="Object Detection Confidence Threshold"
                )
                init_btn = gr.Button("Initialize Batch Processing", variant="primary", elem_classes="btn-primary")
                status_box = gr.Markdown("No batch loaded. Please upload images to begin.", elem_id="status-box")
                
            # Image Preview Card
            with gr.Card():
                gr.Markdown("### 2. Live Verification & Bounding Box Viewer")
                image_viewer = gr.Image(label="Annotated Bounding Box Display", type="numpy", interactive=False)
                
                with gr.Row():
                    prev_btn = gr.Button("◀ Previous Image", variant="secondary", elem_classes="btn-secondary")
                    next_btn = gr.Button("Next Image ▶", variant="secondary", elem_classes="btn-secondary")
                    
        with gr.Column(scale=5):
            # Metadata Review Card
            with gr.Card():
                gr.Markdown("### 3. Interactive Review & Qualitative Labeling")
                tags_editor = gr.Textbox(
                    label="Qualitative Database Labels (Comma-separated, edit freely)",
                    placeholder="Enter keywords or edit auto-extracted categories...",
                    interactive=True
                )
                gr.Markdown(
                    "*The app automatically counts human faces (using Haar-Cascades) and suggests categories (using MobileNet). Feel free to delete, modify, or add qualitative tags for this image above.*"
                )
                
            # Database Sheet Card
            with gr.Card():
                gr.Markdown("### 4. Compiled Database Spreadsheet")
                db_table = gr.Dataframe(
                    headers=["Filename", "Auto-detected Tags", "Face Count", "Final Labels (Edited)"],
                    datatype=["str", "str", "number", "str"],
                    label="Live Database Spreadsheet View",
                    interactive=False,
                    wrap=True,
                    elem_classes="dataframe-container"
                )
                
                export_btn = gr.Button("📊 Compile and Download Finalized CSV Database", variant="primary", elem_classes="btn-primary")
                csv_download = gr.File(label="Download Labeled CSV File", interactive=False)

    # Initialize batch callback
    init_btn.click(
        fn=initialize_batch,
        inputs=[files_input, conf_slider],
        outputs=[image_paths_state, current_idx_state, status_box, image_viewer, db_table, tags_editor, temp_dir_state]
    )
    
    # Navigation callbacks
    prev_btn.click(
        fn=save_and_navigate,
        inputs=[gr.State(-1), current_idx_state, image_paths_state, db_table, tags_editor, conf_slider],
        outputs=[current_idx_state, status_box, image_viewer, db_table, tags_editor]
    )
    
    next_btn.click(
        fn=save_and_navigate,
        inputs=[gr.State(1), current_idx_state, image_paths_state, db_table, tags_editor, conf_slider],
        outputs=[current_idx_state, status_box, image_viewer, db_table, tags_editor]
    )
    
    # Export CSV callback
    export_btn.click(
        fn=export_csv,
        inputs=[db_table],
        outputs=[csv_download]
    )

if __name__ == "__main__":
    demo.launch()