import gradio as gr
from gradio_bbox_annotator import BBoxAnnotator
import json
import os
from pathlib import Path
from PIL import Image
from io import BytesIO
import tempfile

# Define categories and their limits
CATEGORY_LIMITS = {
    "advertisement": 1,  # Maximum 1 advertisement annotation per image
    "text": 2          # Maximum 2 text annotations per image
}
CATEGORIES = list(CATEGORY_LIMITS.keys())
MAX_SIZE = [1024, 1024]  # Maximum width and height for resized images

class AnnotationManager:
    def __init__(self):
        self.annotations = {}
        self.temp_dir = tempfile.mkdtemp()  # Create temporary directory for resized images
    
    def resize_image(self, image_path):
        """Resize image to maximum dimensions while maintaining aspect ratio"""
        try:
            # Read and resize image
            with open(image_path, "rb") as f:
                img = Image.open(BytesIO(f.read()))
                img.thumbnail(MAX_SIZE, Image.Resampling.LANCZOS)
            
            # Save resized image to temporary file
            filename = os.path.basename(image_path)
            temp_path = os.path.join(self.temp_dir, f"resized_{filename}")
            img.save(temp_path)
            
            return temp_path
        except Exception as e:
            raise ValueError(f"Error processing image: {str(e)}")
    
    def process_image_upload(self, image_path):
        """Process uploaded image and return path to resized version"""
        if not image_path:
            return None
        return self.resize_image(image_path)
    
    def validate_annotations(self, bbox_data):
        """Validate the annotation data and return (is_valid, error_message)"""
        if not bbox_data or not isinstance(bbox_data, tuple):
            return False, "No image or annotations provided"
        
        image_path, annotations = bbox_data
        if not isinstance(image_path, str):
            return False, "Invalid image format"
            
        if not annotations:
            return False, "No annotations drawn"
        
        # Count annotations per category
        category_counts = {cat: 0 for cat in CATEGORIES}
        for ann in annotations:
            if len(ann) != 5:
                return False, "Invalid annotation format"
            y1, y2, x1, x2, label = ann
            
            # Validate coordinates
            if any(not isinstance(coord, (int, float)) for coord in [y1, y2, x1, x2]):
                return False, "Invalid coordinate values"
            
            # Validate label
            if not label or label not in CATEGORIES:
                return False, f"Invalid or missing label. Must be one of: {', '.join(CATEGORIES)}"
            
            # Count this annotation
            category_counts[label] += 1
        
        # Check category limits
        for category, count in category_counts.items():
            limit = CATEGORY_LIMITS[category]
            if count > limit:
                return False, f"Too many {category} annotations. Maximum allowed: {limit}"
            
        return True, ""
    
    def add_annotation(self, bbox_data):
        """Add or update annotations for an image"""
        is_valid, error_msg = self.validate_annotations(bbox_data)
        if not is_valid:
            return self.get_json_annotations(), f"❌ Error: {error_msg}"
            
        image_path, annotations = bbox_data
        # Use original filename (remove 'resized_' prefix)
        filename = os.path.basename(image_path)
        if filename.startswith("resized_"):
            filename = filename[8:]
            
        formatted_annotations = []
        for ann in annotations:
            y1, y2, x1, x2, label = ann
            formatted_annotations.append({
                "annotation": [y1, y2, x1, x2],
                "label": label
            })
        self.annotations[filename] = formatted_annotations
        
        # Count annotations by type
        counts = {cat: sum(1 for ann in annotations if ann[4] == cat) for cat in CATEGORIES}
        counts_str = ", ".join(f"{count} {cat}" for cat, count in counts.items())
        success_msg = f"✅ Successfully saved for {filename}: {counts_str}"
        
        return self.get_json_annotations(), success_msg
    
    def get_json_annotations(self):
        """Get all annotations as formatted JSON string"""
        return json.dumps(self.annotations, indent=2)
    
    def clear_annotations(self):
        """Clear all annotations"""
        self.annotations = {}
        return "", "🗑️ All annotations cleared"

def create_interface():
    annotation_mgr = AnnotationManager()
    
    with gr.Blocks() as demo:
        gr.Markdown(f"""
        # Advertisement and Text Annotation Tool
        
        **Instructions:**
        1. Upload an image (will be automatically resized to max {MAX_SIZE[0]}x{MAX_SIZE[1]})
        2. Draw bounding boxes and select the appropriate label
        3. Click 'Save Annotations' to add to the collection
        4. Repeat for all images
        5. Copy the combined JSON when finished
        
        **Annotation Limits per Image:**
        - advertisement: Maximum 1 annotation
        - text: Maximum 2 annotations
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                bbox_input = BBoxAnnotator(
                    show_label=True,
                    label="Draw Bounding Boxes",
                    show_download_button=True,
                    interactive=True,
                    categories=CATEGORIES
                )
            
            with gr.Column(scale=1):
                json_output = gr.TextArea(
                    label="Combined Annotations JSON",
                    interactive=True,
                    lines=15,
                    show_copy_button=True
                )
                
        with gr.Row():
            save_btn = gr.Button("Save Current Image Annotations", variant="primary")
            clear_btn = gr.Button("Clear All Annotations", variant="secondary")

        # Add status message
        status_msg = gr.Markdown(label="Status")
        
        # Event handlers
        def update_image(image_path):
            if not image_path:
                return None
            try:
                resized_path = annotation_mgr.process_image_upload(image_path)
                return resized_path
            except Exception as e:
                return None
        
        # Handle image upload and resizing
        bbox_input.upload(
            fn=update_image,
            inputs=[bbox_input],
            outputs=[bbox_input]
        )
        
        save_btn.click(
            fn=annotation_mgr.add_annotation,
            inputs=[bbox_input],
            outputs=[json_output, status_msg]
        )
        
        clear_btn.click(
            fn=annotation_mgr.clear_annotations,
            inputs=[],
            outputs=[json_output, status_msg]
        )
        
    return demo

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