from PIL import Image, ImageDraw, ImageFilter
import math
import random
import io
import base64

class EndpointHandler:
    def __init__(self, path=""):
        """Initialize DiffSketchEdit handler for Hugging Face Inference API"""
        pass

    def __call__(self, data):
        """Edit sketch based on text prompt"""
        # Extract prompt
        inputs = data.get("inputs", "")
        if isinstance(inputs, dict):
            prompt = inputs.get("prompt", inputs.get("text", ""))
            input_image_data = inputs.get("input_image", None)
            edit_type = inputs.get("edit_type", "refine")
        else:
            prompt = str(inputs)
            input_image_data = data.get("input_image", None)
            edit_type = data.get("edit_type", "refine")
        
        if not prompt:
            prompt = "edit sketch"
        
        # Handle input image if provided
        input_image = None
        if input_image_data:
            try:
                if isinstance(input_image_data, str):
                    # Base64 encoded image
                    img_data = base64.b64decode(input_image_data)
                    input_image = Image.open(io.BytesIO(img_data)).convert('RGB')
                    input_image = input_image.resize((224, 224))
            except Exception:
                input_image = None
        
        # Generate/edit sketch
        image = self.edit_sketch(prompt, input_image, edit_type)
        
        # Return PIL Image directly for HF Inference API
        return image

    def edit_sketch(self, prompt, input_image=None, edit_type="refine"):
        """Edit or create a sketch based on the prompt"""
        
        if input_image is None:
            # Create initial sketch
            img = self._create_initial_sketch(prompt)
        else:
            img = input_image.copy()
        
        # Apply editing based on type
        if edit_type == "refine":
            img = self._refine_sketch(img, prompt)
        elif edit_type == "style_transfer":
            img = self._apply_style_transfer(img, prompt)
        elif edit_type == "add_details":
            img = self._add_details(img, prompt)
        elif edit_type == "color":
            img = self._apply_coloring(img, prompt)
        else:
            # Default refinement
            img = self._refine_sketch(img, prompt)
        
        return img

    def _create_initial_sketch(self, prompt):
        """Create initial sketch based on prompt"""
        img = Image.new('RGB', (224, 224), 'white')
        draw = ImageDraw.Draw(img)
        
        prompt_lower = prompt.lower()
        
        if any(word in prompt_lower for word in ['house', 'building', 'home']):
            self._draw_house_sketch(draw)
        elif any(word in prompt_lower for word in ['tree', 'plant', 'nature']):
            self._draw_tree_sketch(draw)
        elif any(word in prompt_lower for word in ['face', 'portrait', 'person']):
            self._draw_face_sketch(draw)
        elif any(word in prompt_lower for word in ['car', 'vehicle']):
            self._draw_car_sketch(draw)
        elif any(word in prompt_lower for word in ['flower', 'bloom']):
            self._draw_flower_sketch(draw)
        else:
            self._draw_abstract_sketch(draw, prompt)
        
        return img

    def _draw_house_sketch(self, draw):
        """Draw basic house sketch"""
        # Base
        draw.rectangle([50, 120, 174, 180], outline='black', width=2)
        # Roof
        draw.polygon([(50, 120), (112, 80), (174, 120)], outline='black', width=2)
        # Door
        draw.rectangle([100, 150, 124, 180], outline='black', width=2)
        # Windows
        draw.rectangle([70, 140, 90, 160], outline='black', width=2)
        draw.rectangle([134, 140, 154, 160], outline='black', width=2)

    def _draw_tree_sketch(self, draw):
        """Draw basic tree sketch"""
        # Trunk
        draw.rectangle([105, 140, 119, 200], fill='brown', outline='black', width=2)
        # Crown
        draw.ellipse([70, 80, 154, 150], outline='green', width=3)
        # Branches
        for angle in range(0, 360, 60):
            x = 112 + 25 * math.cos(math.radians(angle))
            y = 115 + 25 * math.sin(math.radians(angle))
            draw.line([112, 115, x, y], fill='brown', width=2)

    def _draw_face_sketch(self, draw):
        """Draw basic face sketch"""
        center = (112, 112)
        # Head outline
        draw.ellipse([center[0]-40, center[1]-50, center[0]+40, center[1]+30], 
                    outline='black', width=2)
        # Eyes
        draw.ellipse([center[0]-20, center[1]-20, center[0]-10, center[1]-10], 
                    outline='black', width=2)
        draw.ellipse([center[0]+10, center[1]-20, center[0]+20, center[1]-10], 
                    outline='black', width=2)
        # Nose
        draw.line([center[0], center[1]-5, center[0]-3, center[1]+5], fill='black', width=2)
        # Mouth
        draw.arc([center[0]-15, center[1]+5, center[0]+15, center[1]+20], 0, 180, fill='black', width=2)

    def _draw_car_sketch(self, draw):
        """Draw basic car sketch"""
        # Body
        draw.rectangle([50, 120, 174, 160], outline='black', width=2)
        # Roof
        draw.rectangle([70, 100, 154, 120], outline='black', width=2)
        # Wheels
        draw.ellipse([60, 150, 80, 170], outline='black', width=2)
        draw.ellipse([144, 150, 164, 170], outline='black', width=2)
        # Windows
        draw.rectangle([75, 105, 100, 115], outline='black', width=1)
        draw.rectangle([124, 105, 149, 115], outline='black', width=1)

    def _draw_flower_sketch(self, draw):
        """Draw basic flower sketch"""
        center = (112, 112)
        # Stem
        draw.line([center[0], center[1]+20, center[0], 200], fill='green', width=4)
        # Petals
        for angle in range(0, 360, 45):
            x = center[0] + 25 * math.cos(math.radians(angle))
            y = center[1] + 25 * math.sin(math.radians(angle))
            draw.ellipse([x-8, y-15, x+8, y+5], outline='black', width=2)
        # Center
        draw.ellipse([center[0]-8, center[1]-8, center[0]+8, center[1]+8], 
                    outline='black', width=2)

    def _draw_abstract_sketch(self, draw, prompt):
        """Draw abstract sketch"""
        prompt_hash = hash(prompt) % 100
        
        for i in range(5):
            x1 = (i * 40 + prompt_hash) % 180 + 22
            y1 = (i * 30 + prompt_hash) % 160 + 32
            x2 = x1 + 40 + (i * 10) % 30
            y2 = y1 + 60 + (i * 15) % 40
            draw.ellipse([x1, y1, x2, y2], outline='black', width=2)

    def _refine_sketch(self, img, prompt):
        """Refine the sketch with enhanced details"""
        # Apply sharpening
        img = img.filter(ImageFilter.SHARPEN)
        
        draw = ImageDraw.Draw(img)
        
        # Add refinement details based on prompt
        prompt_lower = prompt.lower()
        
        if "house" in prompt_lower:
            # Add roof tiles
            for y in range(80, 120, 5):
                for x in range(50 + (y-80)//2, 174 - (y-80)//2, 10):
                    draw.line([x, y, x+5, y], fill='gray', width=1)
            # Add door handle
            draw.ellipse([120, 164, 122, 166], fill='black')
            
        elif "tree" in prompt_lower:
            # Add leaf details
            for i in range(10):
                x = 70 + random.randint(0, 84)
                y = 80 + random.randint(0, 70)
                draw.ellipse([x-2, y-2, x+2, y+2], fill='green')
                
        elif "face" in prompt_lower:
            # Add facial details
            center = (112, 112)
            # Eyebrows
            draw.arc([center[0]-25, center[1]-35, center[0]-5, center[1]-25], 0, 180, fill='black', width=2)
            draw.arc([center[0]+5, center[1]-35, center[0]+25, center[1]-25], 0, 180, fill='black', width=2)
            # Pupils
            draw.ellipse([center[0]-17, center[1]-17, center[0]-13, center[1]-13], fill='black')
            draw.ellipse([center[0]+13, center[1]-17, center[0]+17, center[1]-13], fill='black')
        
        return img

    def _apply_style_transfer(self, img, prompt):
        """Apply style transfer effects"""
        prompt_lower = prompt.lower()
        
        if "cartoon" in prompt_lower:
            # Cartoon style - enhance colors and add outlines
            img = img.filter(ImageFilter.EDGE_ENHANCE_MORE)
            # Increase saturation
            from PIL import ImageEnhance
            enhancer = ImageEnhance.Color(img)
            img = enhancer.enhance(1.5)
            
        elif "realistic" in prompt_lower:
            # Realistic style - add shading and texture
            img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
            # Add subtle noise for texture
            import numpy as np
            pixels = np.array(img)
            noise = np.random.normal(0, 3, pixels.shape)
            pixels = np.clip(pixels + noise, 0, 255).astype(np.uint8)
            img = Image.fromarray(pixels)
            
        elif "watercolor" in prompt_lower:
            # Watercolor style - soft edges and bleeding
            img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
            from PIL import ImageEnhance
            enhancer = ImageEnhance.Color(img)
            img = enhancer.enhance(0.8)
            
        return img

    def _add_details(self, img, prompt):
        """Add contextual details to the sketch"""
        draw = ImageDraw.Draw(img)
        prompt_lower = prompt.lower()
        
        if "landscape" in prompt_lower:
            # Add clouds
            for i in range(3):
                x = 40 + i * 60
                y = 30 + i * 10
                draw.ellipse([x, y, x+30, y+15], fill='lightgray', outline='gray')
            # Add birds
            for i in range(2):
                x = 60 + i * 80
                y = 50 + i * 5
                draw.arc([x, y, x+10, y+5], 0, 180, fill='black', width=1)
                
        elif "portrait" in prompt_lower:
            # Add hair details
            center = (112, 112)
            for i in range(15):
                x = center[0] + random.randint(-50, 50)
                y = center[1] - 60 + random.randint(0, 30)
                draw.line([x, y, x + random.randint(-5, 5), y + random.randint(10, 30)], 
                         fill='brown', width=1)
                         
        elif "building" in prompt_lower:
            # Add architectural details
            # Window frames
            draw.rectangle([69, 139, 91, 161], outline='black', width=1)
            draw.rectangle([133, 139, 155, 161], outline='black', width=1)
            # Chimney
            draw.rectangle([140, 70, 150, 90], outline='black', width=2)
            
        return img

    def _apply_coloring(self, img, prompt):
        """Apply coloring to the sketch"""
        # Convert to RGBA for transparency effects
        img = img.convert('RGBA')
        
        # Create color overlay
        color_overlay = Image.new('RGBA', img.size, (0, 0, 0, 0))
        draw = ImageDraw.Draw(color_overlay)
        
        prompt_lower = prompt.lower()
        
        if "sunset" in prompt_lower:
            # Sunset colors
            draw.rectangle([0, 0, 224, 112], fill=(255, 200, 100, 50))
            draw.rectangle([0, 112, 224, 224], fill=(255, 150, 50, 30))
            
        elif "nature" in prompt_lower:
            # Nature colors
            draw.rectangle([0, 140, 224, 224], fill=(100, 200, 100, 40))  # Green ground
            draw.rectangle([0, 0, 224, 140], fill=(150, 200, 255, 30))    # Blue sky
            
        elif "warm" in prompt_lower:
            # Warm color palette
            draw.rectangle([0, 0, 224, 224], fill=(255, 200, 150, 25))
            
        elif "cool" in prompt_lower:
            # Cool color palette
            draw.rectangle([0, 0, 224, 224], fill=(150, 200, 255, 25))
        
        # Blend with original
        img = Image.alpha_composite(img, color_overlay)
        img = img.convert('RGB')
        
        return img