from PIL import Image, ImageDraw
from typing import Dict, Optional, Tuple, List


class SelectionArea:
    """Class to handle selected areas on images (rectangle, circle, or lasso)."""
    
    def __init__(self):
        self.left: Optional[int] = None
        self.top: Optional[int] = None
        self.right: Optional[int] = None
        self.bottom: Optional[int] = None
        self.is_active: bool = False
        self.shape: str = "rectangle"  # "rectangle", "circle", or "lasso"
        self.lasso_path: List[Tuple[int, int]] = []  # List of (x, y) points for lasso selection
    
    def set_coordinates(self, left: int, top: int, right: int, bottom: int) -> None:
        """Set the selection coordinates."""
        self.left = left
        self.top = top
        self.right = right
        self.bottom = bottom
        self.is_active = True
    
    def clear(self) -> None:
        """Clear the selection."""
        self.left = None
        self.top = None
        self.right = None
        self.bottom = None
        self.is_active = False
        self.lasso_path = []
        self.shape = self.shape  # keep last used shape

    def set_shape(self, shape: str) -> None:
        """Set the selection shape: 'rectangle', 'circle', or 'lasso'."""
        if shape in ("rectangle", "circle", "lasso"):
            self.shape = shape
    
    def set_lasso_path(self, path: List[Tuple[int, int]]) -> None:
        """Set the lasso selection path."""
        if not path:
            return
        
        self.lasso_path = path
        self.shape = "lasso"
        self.is_active = True
        
        # Calculate bounding box for lasso
        if path:
            xs = [p[0] for p in path]
            ys = [p[1] for p in path]
            self.left = min(xs)
            self.top = min(ys)
            self.right = max(xs)
            self.bottom = max(ys)
    
    def get_coordinates(self) -> Optional[Tuple[int, int, int, int]]:
        """Get the selection coordinates as a tuple."""
        if self.is_active and all(coord is not None for coord in [self.left, self.top, self.right, self.bottom]):
            return (self.left, self.top, self.right, self.bottom)
        return None
    
    def get_dict(self) -> Optional[Dict[str, int]]:
        """Get the selection coordinates as a dictionary."""
        coords = self.get_coordinates()
        if coords:
            return {
                "left": coords[0],
                "top": coords[1], 
                "right": coords[2],
                "bottom": coords[3]
            }
        return None
    
    def is_valid(self) -> bool:
        """Check if the selection is valid (has area > 0)."""
        if self.shape == "lasso":
            return len(self.lasso_path) >= 3  # Need at least 3 points for a polygon
        coords = self.get_coordinates()
        if coords:
            left, top, right, bottom = coords
            return right > left and bottom > top
        return False
    
    def crop_image(self, image: Image.Image) -> Optional[Image.Image]:
        """Crop the given image to the selection area."""
        if not self.is_valid():
            return None
        
        if self.shape == "lasso":
            # For lasso, crop to bounding box and apply mask
            coords = self.get_coordinates()
            if not coords:
                return None
            left, top, right, bottom = coords
            
            # Crop to bounding box
            cropped = image.crop(coords)
            
            # Create mask for lasso path
            width = right - left
            height = bottom - top
            mask = Image.new("L", (width, height), 0)
            draw = ImageDraw.Draw(mask)
            
            # Adjust lasso path to be relative to bounding box
            adjusted_path = [(x - left, y - top) for x, y in self.lasso_path]
            
            # Draw filled polygon
            if len(adjusted_path) >= 3:
                draw.polygon(adjusted_path, fill=255)
            
            # Apply mask
            result = Image.new("RGBA", (width, height), (0, 0, 0, 0))
            if cropped.mode == "RGBA":
                result.paste(cropped, (0, 0), mask)
            else:
                result.paste(cropped, (0, 0))
                result.putalpha(mask)
            
            return result.convert(image.mode)
        
        coords = self.get_coordinates()
        if coords and self.is_valid():
            # For rectangle and circle, return the bounding box crop
            return image.crop(coords)
        return None
    
    def apply_to_image(self, image: Image.Image, processed_selection: Image.Image) -> Image.Image:
        """Apply the processed selection back to the original image."""
        coords = self.get_coordinates()
        if coords and self.is_valid():
            # Create a copy to avoid modifying the original
            result_image = image.copy()
            width = self.right - self.left
            height = self.bottom - self.top
            
            if self.shape == "circle":
                # Create circular mask for blending only inside the circle
                mask = Image.new("L", (width, height), 0)
                draw = ImageDraw.Draw(mask)
                draw.ellipse((0, 0, width, height), fill=255)
                result_image.paste(processed_selection, (self.left, self.top), mask)
            elif self.shape == "lasso":
                # Create lasso mask for blending only inside the lasso path
                mask = Image.new("L", (width, height), 0)
                draw = ImageDraw.Draw(mask)
                # Adjust lasso path to be relative to bounding box
                adjusted_path = [(x - self.left, y - self.top) for x, y in self.lasso_path]
                if len(adjusted_path) >= 3:
                    draw.polygon(adjusted_path, fill=255)
                result_image.paste(processed_selection, (self.left, self.top), mask)
            else:
                result_image.paste(processed_selection, (self.left, self.top))
            return result_image
        return image
