from model import DesignModel
from PIL import Image
import numpy as np
import random
from typing import List
import traceback

class MockDesignModel(DesignModel):
    def __init__(self):
        super().__init__()
        # Define base colors with proper alpha values
        self.base_colors = [
            (255, 0, 0),      # Red
            (0, 255, 0),      # Green
            (0, 0, 255),      # Blue
            (255, 255, 0),    # Yellow
            (255, 0, 255),    # Magenta
            (0, 255, 255),    # Cyan
            (128, 0, 0),      # Maroon
            (0, 128, 0),      # Dark Green
            (0, 0, 128),      # Navy
        ]
        # Add test-specific attributes
        self.seed = 323*111
        self.neg_prompt = "window, door, low resolution, banner, logo, watermark, text"
        self.additional_quality_suffix = "interior design, 4K, high resolution"

    def apply_tint(self, img_array: np.ndarray, color: tuple) -> np.ndarray:
        """Apply a color tint to an image array"""
        # Create tint array
        tint = np.array(color, dtype=np.float32) / 255.0
        
        # Apply tint with alpha blending
        alpha = 0.3  # 30% tint strength
        tinted = img_array * (1 - alpha) + (img_array * tint) * alpha
        
        # Ensure values are within valid range
        return np.clip(tinted, 0, 255).astype(np.uint8)

    def generate_design(self, image: Image.Image, num_variations: int = 1, **kwargs) -> List[np.ndarray]:
        """Generate multiple variations of the input image with different color tints"""
        try:
            print(f"Starting generation of {num_variations} variations")
            
            # Convert image to numpy array once
            img_array = np.array(image.convert('RGB'))
            
            # Generate base colors for all variations
            colors_needed = max(1, int(num_variations))
            colors = []
            
            # Add base colors first
            colors.extend(self.base_colors)
            
            # Generate additional random colors if needed
            while len(colors) < colors_needed:
                new_color = (
                    random.randint(0, 255),
                    random.randint(0, 255),
                    random.randint(0, 255)
                )
                if new_color not in colors:
                    colors.append(new_color)
            
            # Use only the number of colors we need
            selected_colors = random.sample(colors, colors_needed)
            
            # Generate variations
            variations = []
            for i, color in enumerate(selected_colors):
                # Apply tint to numpy array
                tinted_array = self.apply_tint(img_array.copy(), color)
                variations.append(tinted_array)
                print(f"Created variation {i+1}/{colors_needed}")
            
            print(f"Successfully generated {len(variations)} variations")
            return variations
            
        except Exception as e:
            print(f"Error in generate_design: {e}")
            traceback.print_exc()
            # Return the original image array if there's an error
            return [np.array(image.convert('RGB'))]