from aiforge.core.registry import register
from aiforge.core.context import BaseOperator

@register("blur")
class Blur(BaseOperator):
    """Apply Gaussian Blur to an image using OpenCV."""
    def run(self, data, kernel=5):
        import cv2
        return cv2.GaussianBlur(data, (kernel, kernel), 0)

@register("grayscale")
class Grayscale(BaseOperator):
    """Convert an image to grayscale using OpenCV."""
    def run(self, data):
        import cv2
        return cv2.cvtColor(data, cv2.COLOR_BGR2GRAY)

@register("edge_detect")
class EdgeDetect(BaseOperator):
    """Apply Canny edge detection to an image."""
    def run(self, data, threshold1=100, threshold2=200):
        import cv2
        return cv2.Canny(data, threshold1, threshold2)

@register("resize_image")
class ResizeImage(BaseOperator):
    """Resize an image to a given width and height."""
    def run(self, data, width=224, height=224):
        import cv2
        return cv2.resize(data, (width, height))

@register("threshold_image")
class ThresholdImage(BaseOperator):
    """Apply binary thresholding to an image."""
    def run(self, data, threshold_value=127, max_value=255):
        import cv2
        # Ensure image is grayscale first or handle BGR
        if len(data.shape) == 3:
            data = cv2.cvtColor(data, cv2.COLOR_BGR2GRAY)
        _, thresholded = cv2.threshold(data, threshold_value, max_value, cv2.THRESH_BINARY)
        return thresholded

@register("rotate_image")
class RotateImage(BaseOperator):
    """Rotate an image by a given angle in degrees."""
    def run(self, data, angle=90.0):
        import cv2
        (h, w) = data.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, angle, 1.0)
        return cv2.warpAffine(data, M, (w, h))
