import cv2
import numpy as np
import logging
from pathlib import Path
from ultralytics import YOLO
from exceptions import ModelLoadError, DetectionError, ImageProcessingError

logger = logging.getLogger(__name__)


class MemoryDetector:
    def __init__(self, model_path, confidence_threshold=0.3, image_size=416):
        self.model_path = model_path
        self.confidence_threshold = confidence_threshold
        self.image_size = image_size
        self.model = None
        self._load_model()
    
    def _load_model(self):
        """Load YOLO model with error handling."""
        try:
            if not Path(self.model_path).exists():
                raise FileNotFoundError(f"Model file not found: {self.model_path}")
            
            logger.info(f"Loading model from {self.model_path}")
            self.model = YOLO(self.model_path)
            logger.info("Model loaded successfully")
            
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            raise ModelLoadError(f"Model loading failed: {str(e)}")
    
    def detect_from_bytes(self, image_bytes):
        """Detect memory modules from image bytes."""
        try:
            # Decode image
            nparr = np.frombuffer(image_bytes, np.uint8)
            image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            
            if image is None:
                raise ImageProcessingError("Could not decode image")
            
            return self._perform_detection(image)
            
        except Exception as e:
            logger.error(f"Detection from bytes failed: {e}")
            raise DetectionError(f"Detection failed: {str(e)}")
    
    def detect_from_file(self, file_path):
        """Detect memory modules from image file."""
        try:
            if not Path(file_path).exists():
                raise FileNotFoundError(f"Image file not found: {file_path}")
            
            image = cv2.imread(str(file_path))
            if image is None:
                raise ImageProcessingError("Could not load image file")
            
            return self._perform_detection(image)
            
        except Exception as e:
            logger.error(f"Detection from file failed: {e}")
            raise DetectionError(f"Detection failed: {str(e)}")
    
    def _perform_detection(self, image):
        """Perform detection on image."""
        try:
            logger.info("Running detection")
            
            # Run inference
            results = self.model.predict(
                image,
                imgsz=self.image_size,
                conf=self.confidence_threshold,
                verbose=False
            )
            
            # Extract detections
            detections = self._extract_detections(results[0])
            
            # Create annotated image
            annotated_image = self._draw_boxes(image, detections)
            
            logger.info(f"Detection completed: {len(detections)} objects found")
            
            return {
                'detections': detections,
                'annotated_image': annotated_image,
                'detection_count': len(detections)
            }
            
        except Exception as e:
            logger.error(f"Detection processing failed: {e}")
            raise DetectionError(f"Detection processing failed: {str(e)}")
    
    def _extract_detections(self, result):
        """Extract detection results."""
        detections = []
        
        if result.boxes is not None:
            boxes = result.boxes.xyxy.cpu().numpy()
            confidences = result.boxes.conf.cpu().numpy()
            classes = result.boxes.cls.cpu().numpy() if result.boxes.cls is not None else None
            
            for i, (box, conf) in enumerate(zip(boxes, confidences)):
                detection = {
                    'box': [float(coord) for coord in box],  # [x1, y1, x2, y2]
                    'confidence': float(conf),
                    'class': int(classes[i]) if classes is not None else 0
                }
                detections.append(detection)
        
        return detections
    
    def _draw_boxes(self, image, detections):
        """Draw bounding boxes on image."""
        annotated = image.copy()
        
        for detection in detections:
            box = detection['box']
            confidence = detection['confidence']
            
            # Extract coordinates
            x1, y1, x2, y2 = map(int, box)
            
            # Draw bounding box
            cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 2)
            
            # Draw confidence score
            label = f"Memory: {confidence:.2f}"
            cv2.putText(
                annotated, label, (x1, y1 - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1
            )
        
        return annotated