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

CATEGORIES = ["advertisement", "text"]

class AnnotationManager:
    def __init__(self):
        self.annotations = {}
    
    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"
        
        # Check each annotation
        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)}"
            
        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(), error_msg
            
        image_path, annotations = bbox_data
        filename = os.path.basename(image_path)
        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
        ad_count = sum(1 for ann in annotations if ann[4] == "advertisement")
        text_count = sum(1 for ann in annotations if ann[4] == "text")
        success_msg = f"Successfully saved for {filename}: {ad_count} advertisements, {text_count} text regions"
        
        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("""
        # 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. Copy the combined 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
                )
            
            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 with different styling based on message type
        status_msg = gr.Markdown(label="Status")
        
        # Event handlers
        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()