import logging
import numpy as np
from typing import List, Dict, Any
from modules.detectors.base_detector import ObjectDetector
from modules.config import UI_MODEL_PATH, PYTORCH_ENABLE_MPS_FALLBACK

logger = logging.getLogger(__name__)

class RectangleDetector(ObjectDetector):
    """Detector for rectangular UI elements like buttons, input fields, etc."""
    
    def __init__(self):
        self.model = None
        # self.ui_classes = ['button', 'card', 'field', 'heading', 'icon', 'image', 'link', 'paragraph', 'text']
        self.ui_classes = ['card', 'field', 'icon', 'image', 'paragraph']
    
    async def initialize(self):
        """Initialize the YOLO model for UI element detection."""
        from ultralytics import YOLO
        import torch
        import os
        
        # Check for available hardware acceleration
        if torch.cuda.is_available():
            self.device = "cuda:0"
            logger.info("CUDA is available, using GPU acceleration for UI detection")
        elif hasattr(torch, 'mps') and torch.mps.is_available():
            self.device = "mps"
            logger.info("MPS is available, using Apple Silicon acceleration for UI detection")
            # Enable MPS fallback for operations not supported by MPS
            os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = PYTORCH_ENABLE_MPS_FALLBACK
        else:
            self.device = "cpu"
            logger.warning("No hardware acceleration available, falling back to CPU for UI detection")
        
        # Initialize UI element detection model with custom-trained weights
        model_path = UI_MODEL_PATH
        
        if os.path.exists(model_path):
            self.model = YOLO(model_path)
            logger.info(f"UI element detector initialized from {model_path} on {self.device}")
        else:
            logger.error(f"UI model file {model_path} not found")
            raise FileNotFoundError(f"UI model file {model_path} not found")
    
    def detect(self, image: np.ndarray) -> List[Dict]:
        """Detect UI elements in an image."""
        if self.model is None:
            raise RuntimeError("UI detection model not initialized")
        
        # Ensure image has 3 channels (RGB)
        if len(image.shape) == 3 and image.shape[2] == 4:  # If RGBA, convert to RGB
            image = image[:, :, :3]
            
        # Run inference
        results = self.model(image, conf=0.25, device=self.device)
        
        # Process and format results
        detected_elements = []
        for result in results:
            boxes = result.boxes
            for box in boxes:
                # Get bounding box coordinates
                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
                
                # Get class name
                cls_id = int(box.cls[0].item())
                cls_name = result.names[cls_id]
                confidence = float(box.conf[0].item())
                
                # Only include UI elements
                # Note: This is a placeholder - with a proper UI element model,
                # this filtering wouldn't be necessary
                if cls_name.lower() in self.ui_classes:
                    detected_elements.append({
                        "label": cls_name.lower(),
                        "confidence": confidence,
                        "bbox": [x1, y1, x2, y2]
                    })
        
        return detected_elements