utils/validators.py

import re
from typing import Optional, Tuple
from PIL import Image
from io import BytesIO
from config import Config


class ValidationError(Exception):
    """Validation error"""
    pass


def validate_api_key(api_key: str) -> bool:
    """Validate API key format"""
    if not api_key:
        return False

    # Basic format check: length and characters
    if len(api_key) < 20 or not re.match(r'^[A-Za-z0-9_\-]+$', api_key):
        return False

    return True


def validate_prompt(prompt: str, max_length: int = 2000) -> Tuple[bool, Optional[str]]:
    """
    Validate prompt

    Returns:
        (is_valid, error_message)
    """
    if not prompt or not prompt.strip():
        return False, "Prompt cannot be empty"

    if len(prompt) > max_length:
        return False, f"Prompt length cannot exceed {max_length} characters"

    # Check for potentially harmful content keywords
    harmful_keywords = [
        'violence', 'gore', 'explicit', 'nsfw', 'nude', 'sexual',
        '暴力', '色情', '裸体', '性', '血腥'
    ]

    prompt_lower = prompt.lower()
    for keyword in harmful_keywords:
        if keyword in prompt_lower:
            return False, f"Prompt contains inappropriate content: {keyword}"

    return True, None


def validate_image(image_file) -> Tuple[bool, Optional[str], Optional[Image.Image]]:
    """
    Validate uploaded image

    Returns:
        (is_valid, error_message, processed_image)
    """
    try:
        # Check file size
        if hasattr(image_file, 'size') and image_file.size > Config.MAX_FILE_SIZE_MB * 1024 * 1024:
            return False, f"Image size cannot exceed {Config.MAX_FILE_SIZE_MB}MB", None

        # Try to open image
        image = Image.open(image_file)

        # Check image format
        if image.format.lower() not in ['png', 'jpeg', 'jpg']:
            return False, "Only PNG, JPG, JPEG formats are supported", None

        # Check image dimensions
        width, height = image.size

        # Minimum size check
        if width < 64 or height < 64:
            return False, "Image size too small (minimum 64x64 pixels)", None

        # Maximum size check
        max_dimension = 4096
        if width > max_dimension or height > max_dimension:
            return False, f"Image size too large (maximum {max_dimension}x{max_dimension} pixels)", None

        # Aspect ratio check (extreme ratios may not work well)
        ratio = max(width, height) / min(width, height)
        if ratio > 10:
            return False, "Image aspect ratio too extreme", None

        # If RGBA, convert to RGB
        if image.mode == 'RGBA':
            rgb_image = Image.new('RGB', image.size, (255, 255, 255))
            rgb_image.paste(image, mask=image.split()[-1])
            image = rgb_image
        elif image.mode != 'RGB':
            image = image.convert('RGB')

        return True, None, image

    except Exception as e:
        return False, f"Image processing error: {str(e)}", None


def validate_model_parameters(model_type: str, resolution: str, duration: int,
                            aspect_ratio: str) -> Tuple[bool, Optional[str]]:
    """Validate model parameters"""

    # Validate model type
    valid_models = list(Config.MODEL_OPTIONS.values())
    if model_type not in valid_models:
        return False, f"Unsupported model type: {model_type}"

    # Validate resolution
    if resolution not in Config.RESOLUTIONS:
        return False, f"Unsupported resolution: {resolution}"

    # Validate duration
    if duration not in Config.DURATIONS:
        return False, f"Unsupported duration: {duration}"

    # Validate aspect ratio
    if aspect_ratio not in Config.ASPECT_RATIOS:
        return False, f"Unsupported aspect ratio: {aspect_ratio}"

    return True, None


def sanitize_filename(filename: str) -> str:
    """Clean filename, remove unsafe characters"""
    # Remove path separators and special characters
    sanitized = re.sub(r'[<>:"/\\|?*]', '_', filename)
    # Remove control characters
    sanitized = re.sub(r'[\x00-\x1f\x7f-\x9f]', '', sanitized)
    # Limit length
    if len(sanitized) > 255:
        name, ext = sanitized.rsplit('.', 1) if '.' in sanitized else (sanitized, '')
        sanitized = name[:250] + ('.' + ext if ext else '')

    return sanitized or 'untitled'


def validate_seed(seed: Optional[int]) -> Tuple[bool, Optional[str]]:
    """Validate random seed"""
    if seed is None:
        return True, None

    if not isinstance(seed, int):
        return False, "Seed must be an integer"

    if seed < 0 or seed > 2147483647:
        return False, "Seed value must be between 0 and 2147483647"

    return True, None