import gradio as gr
from gradio_bbox_annotator import BBoxAnnotator
import json
import os
from pathlib import Path

# Define the allowed categories
CATEGORIES = ["advertisement", "text"]

class AnnotationManager:
    def __init__(self):
        self.annotations = {}
    
    def add_annotation(self, bbox_data):
        """Add or update annotations for an image"""
        if bbox_data and isinstance(bbox_data, tuple):
            image_path, annotations = bbox_data
            if isinstance(image_path, str) and annotations:
                filename = os.path.basename(image_path)
                formatted_annotations = []
                for ann in annotations:
                    y1, y2, x1, x2, label = ann
                    # Ensure label is one of the allowed categories
                    if label not in CATEGORIES:
                        continue
                    formatted_annotations.append({
                        "annotation": [y1, y2, x1, x2],
                        "label": label
                    })
                if formatted_annotations:  # Only add if there are valid annotations
                    self.annotations[filename] = formatted_annotations
        return self.get_json_annotations()
    
    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 ""

def create_interface():
    annotation_mgr = AnnotationManager()
    
    with gr.Blocks() as demo:
        gr.Markdown("""
        # Advertisement and Text Annotation Tool
        
        **Instructions:**
        1. Upload an image using the upload button in the annotator
        2. Draw bounding boxes and select either 'advertisement' or 'text' as label
        3. Click 'Save Annotations' to add to the collection
        4. Repeat for all images
        5. Click 'Download JSON' when finished
        
        **Labels:**
        - advertisement: Use for marking advertisement regions
        - text: Use for marking text regions
        """)
        
        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  # Set the predefined categories
                )
            
            with gr.Column(scale=1):
                json_output = gr.TextArea(
                    label="Combined Annotations JSON",
                    interactive=False,
                    lines=15
                )
                
        with gr.Row():
            save_btn = gr.Button("Save Current Image Annotations", variant="primary")
            clear_btn = gr.Button("Clear All Annotations", variant="secondary")
            download_btn = gr.Button("Download Combined JSON", variant="primary")

        # Add status message
        status_msg = gr.Textbox(label="Status", interactive=False)
        
        # Event handlers
        def save_and_update(bbox_data):
            if not bbox_data or not isinstance(bbox_data, tuple):
                return annotation_mgr.get_json_annotations(), "No image or annotations to save"
            
            image_path, annotations = bbox_data
            if not annotations:
                return annotation_mgr.get_json_annotations(), "No annotations drawn"
            
            # Count annotations by type
            ad_count = sum(1 for ann in annotations if ann[4] == "advertisement")
            text_count = sum(1 for ann in annotations if ann[4] == "text")
                
            json_str = annotation_mgr.add_annotation(bbox_data)
            filename = os.path.basename(image_path) if isinstance(image_path, str) else "unknown"
            return json_str, f"Saved for {filename}: {ad_count} advertisements, {text_count} text regions"
        
        save_btn.click(
            fn=save_and_update,
            inputs=[bbox_input],
            outputs=[json_output, status_msg]
        )
        
        clear_btn.click(
            fn=lambda: (annotation_mgr.clear_annotations(), "All annotations cleared"),
            inputs=[],
            outputs=[json_output, status_msg]
        )
        
        download_btn.click(
            fn=lambda x: {"text": x, "filename": "combined_annotations.json"},
            inputs=[json_output],
            outputs=[gr.File()]
        )
        
    return demo

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