from ultralytics import YOLO
from PIL import Image
import gradio as gr
from huggingface_hub import snapshot_download
import os
import shutil

def initialize_model_path():
    is_huggingface = os.environ.get("SPACE_ID") is not None
    if is_huggingface:
        REPO_ID = "mail2kandan/guitar_watch_openvino_model"
        print("Downloading model from Hugging Face...")
        
        # Create a temp directory for the model
        model_dir = os.path.join(os.getcwd(), 'guitar_openvino_model')
        os.makedirs(model_dir, exist_ok=True)
        
        # Download complete model directory
        downloaded_path = snapshot_download(repo_id=REPO_ID)
        
        # Copy files to our model directory
        for ext in ['.xml', '.bin']:
            src = os.path.join(downloaded_path, f'best{ext}')
            dst = os.path.join(model_dir, f'best{ext}')
            if os.path.exists(src):
                shutil.copy2(src, dst)
                print(f"Copied {src} to {dst}")
            else:
                raise Exception(f"Required model file best{ext} not found in repository")
        
        # Verify both files exist in the new directory
        if os.path.exists(os.path.join(model_dir, 'best.xml')) and \
           os.path.exists(os.path.join(model_dir, 'best.bin')):
            print(f"Using model directory: {model_dir}")
            return model_dir
        else:
            raise Exception("Model files not found in target directory after copying")
    else:
        local_path = r'C:\Manikandan\NYP\ITI107\Assignment_draft\output_models\best_openvino_model'
        print(f"Using local model from: {local_path}")
        return local_path

# Initialize model path once
MODEL_PATH = initialize_model_path()

def predict(pilimg, conf_threshold, iou_threshold):
    try:
        if pilimg is None:
            return None
            
        print(f"Loading model from: {MODEL_PATH}")
        print(f"Model path exists: {os.path.exists(MODEL_PATH)}")
        if os.path.isfile(MODEL_PATH):
            print(f"Found these files in directory: {os.listdir(os.path.dirname(MODEL_PATH))}")
        
        model = YOLO(MODEL_PATH, task='detect')
        result = model.predict(pilimg, conf=conf_threshold, iou=iou_threshold)
        img_bgr = result[0].plot()
        return Image.fromarray(img_bgr[..., ::-1])
    except Exception as e:
        print(f"Error during prediction: {str(e)}")
        # Print full error traceback for debugging
        import traceback
        print(traceback.format_exc())
        return None

with gr.Blocks() as demo:
    # Header
    gr.HTML("""
        <div style="text-align: center; margin-bottom: 20px; padding: 20px; background-color: #f7f7f7;">
            <h1 style="font-size: 2.5em; color: #2196F3; margin-bottom: 10px;">Guitar & Watch Detection</h1>
            <p style="font-size: 1.2em; color: #666;">Powered by YOLOv8</p>
        </div>
    """)
    
    # Hero section with image and description side by side
    with gr.Row():
        if os.path.exists("testImages/heroimage.png"):
            with gr.Column(scale=1):
                gr.Image("testImages/heroimage.png", 
                         label="Watch Detection",
                         show_label=False,
                         elem_id="hero-image")
        
        with gr.Column(scale=1):
            gr.Markdown("""
                ### 🎯 About This App
                This application uses Yolo object detection to identify guitars and wristwatches in images. 
                The model is optimized using small model '8s' for faster inference.
                
                ### 📝 How to Use
                1. Upload an image or use one of the examples below
                2. Adjust the confidence and IOU thresholds if needed
                3. View the detection results in real-time
            """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                type="pil",
                label="📸 Upload Image"
            )
            
            gr.Markdown("### ⚙️ Detection Parameters")
            with gr.Column(visible=True):
                conf_threshold = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    step=0.05,
                    value=0.5,
                    label="🎯 Confidence Threshold",
                    info="Higher values mean more confident detections"
                )
                iou_threshold = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    step=0.05,
                    value=0.6,
                    label="🔄 IOU Threshold",
                    info="Higher values mean less box overlap"
                )
            
        with gr.Column():
            output_image = gr.Image(
                type="pil",
                label="🎉 Detection Result"
            )
    
    # Examples section
    gr.Markdown("### 📚 Example Images")
    if os.path.exists("testImages/acoustic_guitar_005.jpg"):
        examples = gr.Examples(
            examples=[
                ["testImages/acoustic_guitar_005.jpg"],
                ["testImages/acoustic_guitar_006.jpg"],
                ["testImages/analog_watch_024.jpg"],
                ["testImages/analog_watch_025.jpg"]
            ],
            inputs=input_image,
            label="Try these examples"
        )
    
    # Footer
    gr.HTML("""
        <div style="text-align: center; margin-top: 30px; padding: 20px; background-color: #f7f7f7;">
            <p style="color: #666; font-size: 0.9em;">© 2025 Guitar & Watch Detection - 9027941P - Manikandan Sadhasivam</p>
        </div>
    """)
    
    # Set up the prediction functionality
    input_image.change(
        fn=predict,
        inputs=[input_image, conf_threshold, iou_threshold],
        outputs=output_image
    )

    # Add custom CSS to resize the hero image
    demo.css = """
    #hero-image img {
        width: 560px;
        height: auto;
    }
    """

if __name__ == "__main__":
    is_huggingface = os.environ.get("SPACE_ID") is not None
    demo.launch(share=not is_huggingface)