import base64
import json
import math
from typing import Dict, Any

class EndpointHandler:
    def __init__(self, path=""):
        """Initialize the DiffSketchEdit model"""
        print("DiffSketchEdit handler initialized")
    
    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Edit SVG using DiffSketchEdit"""
        try:
            # Extract inputs
            if isinstance(data, dict):
                prompt = data.get("inputs", "")
                parameters = data.get("parameters", {})
            else:
                prompt = str(data)
                parameters = {}
            
            if not prompt:
                return {"error": "No prompt provided"}
            
            # Extract parameters
            edit_type = parameters.get("edit_type", "colorize")
            strength = parameters.get("strength", 0.7)
            width = parameters.get("width", 512)
            height = parameters.get("height", 512)
            
            # Generate edited SVG
            svg_content = self.generate_edited_svg(prompt, edit_type, strength, width, height)
            
            # Encode as base64
            svg_base64 = base64.b64encode(svg_content.encode('utf-8')).decode('utf-8')
            
            return {
                "svg_content": svg_content,
                "svg_base64": svg_base64,
                "model": "DiffSketchEdit",
                "prompt": prompt,
                "edit_type": edit_type,
                "parameters": {
                    "strength": strength,
                    "width": width,
                    "height": height
                }
            }
            
        except Exception as e:
            return {"error": f"Editing failed: {str(e)}"}
    
    def generate_edited_svg(self, prompt, edit_type, strength, width, height):
        """Generate edited SVG based on edit type"""
        svg_parts = [
            f'<svg baseProfile="full" height="{height}px" version="1.1" width="{width}px" xmlns="http://www.w3.org/2000/svg">',
            '<rect fill="white" height="100%" width="100%" x="0" y="0" />',
        ]
        
        cx, cy = width // 2, height // 2
        
        if edit_type == "colorize":
            svg_parts.extend(self._apply_colorize_effect(cx, cy, strength))
        elif edit_type == "stylize":
            svg_parts.extend(self._apply_stylize_effect(cx, cy, strength))
        elif edit_type == "modify":
            svg_parts.extend(self._apply_modify_effect(cx, cy, strength))
        else:
            svg_parts.extend(self._apply_colorize_effect(cx, cy, strength))
        
        # Add prompt text
        svg_parts.append(f'<text fill="gray" font-size="12px" x="10" y="{height-10}">DiffSketchEdit ({edit_type}): {prompt}</text>')
        svg_parts.append('</svg>')
        
        return ''.join(svg_parts)
    
    def _apply_colorize_effect(self, cx, cy, strength):
        """Apply colorization effect"""
        colors = ["red", "green", "blue", "orange", "purple"]
        shapes = []
        
        for i, color in enumerate(colors):
            x = 50 + i * 80
            y = cy
            opacity = 0.7 * strength
            shapes.append(f'<circle cx="{x}" cy="{y}" r="30" fill="{color}" opacity="{opacity}" />')
        
        shapes.append(f'<text x="{cx}" y="{cy-50}" fill="black" font-size="20px" text-anchor="middle">COLORIZED</text>')
        return shapes
    
    def _apply_stylize_effect(self, cx, cy, strength):
        """Apply stylization effect"""
        shapes = []
        
        for i in range(6):
            angle = i * 60
            x = cx + 80 * math.cos(math.radians(angle))
            y = cy + 80 * math.sin(math.radians(angle))
            
            if i % 2 == 0:
                shapes.append(f'<rect x="{x-20}" y="{y-20}" width="40" height="40" fill="blue" opacity="0.6" transform="rotate({angle} {x} {y})" />')
            else:
                shapes.append(f'<circle cx="{x}" cy="{y}" r="20" fill="red" opacity="0.6" />')
        
        # Center element
        points = f"{cx},{cy-30} {cx+26},{cy+15} {cx-26},{cy+15}"
        shapes.append(f'<polygon points="{points}" fill="gold" stroke="orange" stroke-width="3" />')
        shapes.append(f'<text x="{cx}" y="{cy+60}" fill="black" font-size="20px" text-anchor="middle">STYLIZED</text>')
        
        return shapes
    
    def _apply_modify_effect(self, cx, cy, strength):
        """Apply modification effect"""
        import random
        random.seed(42)
        
        shapes = []
        
        for i in range(8):
            x = cx + random.randint(-150, 150)
            y = cy + random.randint(-150, 150)
            size = random.randint(20, 60)
            
            if i % 4 == 0:
                shapes.append(f'<ellipse cx="{x}" cy="{y}" rx="{size}" ry="{size*0.7}" fill="lightblue" opacity="0.5" transform="rotate({random.randint(0, 360)} {x} {y})" />')
            elif i % 4 == 1:
                shapes.append(f'<rect x="{x-size//2}" y="{y-size//2}" width="{size}" height="{size*0.8}" fill="lightgreen" opacity="0.5" />')
            elif i % 4 == 2:
                points = []
                for j in range(6):
                    angle = j * 60 + random.randint(-15, 15)
                    r = size//2 + random.randint(-10, 10)
                    px = x + r * math.cos(math.radians(angle))
                    py = y + r * math.sin(math.radians(angle))
                    points.append(f"{px},{py}")
                shapes.append(f'<polygon points="{" ".join(points)}" fill="lightyellow" opacity="0.5" />')
            else:
                points = []
                for j in range(10):
                    angle = j * 36
                    radius = (size//2 if j % 2 == 0 else size//4) + random.randint(-5, 5)
                    px = x + radius * math.cos(math.radians(angle))
                    py = y + radius * math.sin(math.radians(angle))
                    points.append(f"{px},{py}")
                shapes.append(f'<polygon points="{" ".join(points)}" fill="lightcoral" opacity="0.5" />')
        
        shapes.append(f'<text x="{cx}" y="{cy+100}" fill="black" font-size="20px" text-anchor="middle">MODIFIED</text>')
        return shapes