#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
图像处理组件 - 新架构版本
融合demo_refactored_templates.py和poster_template.py的图像处理功能
确保视觉效果完全一致并具有出色的性能
"""

import os
import gc
import logging
import psutil
from typing import Tuple, List, Dict, Any, Optional, Union
from pathlib import Path
from PIL import Image, ImageDraw, ImageFilter, ImageEnhance
import numpy as np
from .text_renderer import PosterComponent

logger = logging.getLogger(__name__)


class ImageProcessor(PosterComponent):
    """
    图像处理器 - 融合版本
    结合原版的高质量图像处理和优化版的内存管理
    """
    
    def __init__(self, max_memory_percent=80, target_max_pixels=4000000):
        self.logger = logging.getLogger(__name__)
        self.max_memory_percent = max_memory_percent
        self.target_max_pixels = target_max_pixels
        self._loaded_images = {}  # 图像缓存
        self._memory_pressure_threshold = 0.85
        
    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化图像处理器"""
        self.max_memory_percent = config.get('max_memory_percent', self.max_memory_percent)
        self.target_max_pixels = config.get('target_max_pixels', self.target_max_pixels)
        self._memory_pressure_threshold = config.get('memory_pressure_threshold', 0.85)
        
        self.logger.info(f"图像处理器初始化 - 内存限制: {self.max_memory_percent}%, 像素限制: {self.target_max_pixels}")
    
    def process(self, input_data: Any, **kwargs) -> Any:
        """处理图像相关请求"""
        if isinstance(input_data, dict):
            operation = input_data.get('operation')
            params = input_data.get('params', {})
            
            if operation == 'load_image':
                return self.smart_load_image(**params)
            elif operation == 'resize_crop':
                return self.resize_and_crop(**params)
            elif operation == 'enhance':
                return self.enhance_image(**params)
            elif operation == 'apply_blur':
                return self.apply_blur(**params)
            elif operation == 'batch_process':
                return self.batch_process_directory(**params)
                
        return None
    
    def cleanup(self) -> None:
        """清理资源"""
        self.release_loaded_images()
        self.force_garbage_collection()
    
    # === 内存管理功能 (来自poster_template.py) ===
    
    def get_memory_usage(self) -> Tuple[float, float]:
        """
        获取当前内存使用情况 - 与原版完全一致
        
        Returns:
            (使用百分比, 可用内存MB)
        """
        memory = psutil.virtual_memory()
        used_percent = memory.percent
        available_mb = memory.available / 1024 / 1024
        return used_percent, available_mb
    
    def check_memory_pressure(self) -> bool:
        """
        检查内存压力 - 与原版逻辑完全一致
        
        Returns:
            True if memory pressure is high
        """
        used_percent, available_mb = self.get_memory_usage()
        
        # 检查多个条件
        conditions = [
            used_percent > self.max_memory_percent,  # 使用率过高
            available_mb < 500,  # 可用内存小于500MB
            used_percent > self._memory_pressure_threshold * 100  # 超过压力阈值
        ]
        
        if any(conditions):
            self.logger.warning(f"内存压力检测 - 使用率: {used_percent:.1f}%, 可用: {available_mb:.1f}MB")
            return True
            
        return False
    
    def force_garbage_collection(self):
        """
        强制垃圾回收 - 与原版完全一致
        """
        try:
            initial_memory = psutil.virtual_memory().percent
            
            # 多次垃圾回收确保彻底清理
            for _ in range(3):
                gc.collect()
            
            final_memory = psutil.virtual_memory().percent
            memory_freed = initial_memory - final_memory
            
            if memory_freed > 0:
                self.logger.info(f"垃圾回收完成 - 释放了 {memory_freed:.1f}% 内存")
            
        except Exception as e:
            self.logger.error(f"垃圾回收失败: {e}")
    
    def release_loaded_images(self):
        """
        释放已加载的图像 - 与原版完全一致
        """
        try:
            released_count = len(self._loaded_images)
            self._loaded_images.clear()
            
            if released_count > 0:
                self.logger.info(f"释放了 {released_count} 个缓存图像")
                self.force_garbage_collection()
                
        except Exception as e:
            self.logger.error(f"释放图像缓存失败: {e}")
    
    def calculate_resize_factor(self, width: int, height: int) -> float:
        """
        计算调整因子 - 与原版算法完全一致
        
        Args:
            width: 图像宽度
            height: 图像高度
            
        Returns:
            调整因子
        """
        current_pixels = width * height
        
        if current_pixels <= self.target_max_pixels:
            return 1.0
        
        # 计算需要的缩放比例
        resize_factor = (self.target_max_pixels / current_pixels) ** 0.5
        
        # 确保因子在合理范围内
        resize_factor = max(0.1, min(1.0, resize_factor))
        
        # 为了避免过度精确，量化到常见的因子
        common_factors = [0.1, 0.2, 0.25, 0.33, 0.5, 0.67, 0.75, 0.8, 0.9, 1.0]
        resize_factor = min(common_factors, key=lambda x: abs(x - resize_factor))
        
        return resize_factor
    
    def smart_load_image(self, image_path: str, target_size: Optional[Tuple[int, int]] = None) -> Image.Image:
        """
        智能图像加载 - 与原版完全一致
        
        Args:
            image_path: 图像路径
            target_size: 目标尺寸
            
        Returns:
            加载的图像
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图像文件不存在: {image_path}")
        
        # 检查缓存
        cache_key = f"{image_path}_{target_size}"
        if cache_key in self._loaded_images:
            return self._loaded_images[cache_key]
        
        # 检查内存压力
        if self.check_memory_pressure():
            self.logger.info("内存压力过高，使用保守加载模式")
            return self._conservative_load(image_path, target_size)
        
        try:
            # 正常加载
            image = Image.open(image_path)
            
            # 获取原始尺寸
            original_width, original_height = image.size
            self.logger.debug(f"加载图像: {image_path} ({original_width}x{original_height})")
            
            # 计算是否需要预调整尺寸
            resize_factor = self.calculate_resize_factor(original_width, original_height)
            
            if resize_factor < 1.0:
                new_width = int(original_width * resize_factor)
                new_height = int(original_height * resize_factor)
                image = image.resize((new_width, new_height), Image.LANCZOS)
                self.logger.debug(f"预调整图像尺寸: {new_width}x{new_height} (因子: {resize_factor:.2f})")
            
            # 如果指定了目标尺寸，进一步调整
            if target_size:
                image = self.resize_and_crop(image, target_size)
            
            # 确保图像格式
            if image.mode not in ['RGB', 'RGBA']:
                image = image.convert('RGBA' if image.mode == 'P' and 'transparency' in image.info else 'RGB')
            
            # 缓存图像（如果内存允许）
            current_memory = psutil.virtual_memory().percent
            if current_memory < self.max_memory_percent * 0.8:  # 只在内存充足时缓存
                self._loaded_images[cache_key] = image
            
            return image
            
        except Exception as e:
            self.logger.error(f"智能加载图像失败 {image_path}: {e}")
            return self._conservative_load(image_path, target_size)
    
    def _conservative_load(self, image_path: str, target_size: Optional[Tuple[int, int]] = None) -> Image.Image:
        """
        保守的图像加载方式 - 与原版完全一致
        """
        try:
            # 先释放一些内存
            self.release_loaded_images()
            self.force_garbage_collection()
            
            # 使用最小内存方式加载
            with Image.open(image_path) as img:
                # 立即复制并关闭原图像
                image = img.copy()
            
            # 强制调整尺寸以减少内存使用
            width, height = image.size
            resize_factor = min(0.5, self.calculate_resize_factor(width, height))
            
            if resize_factor < 1.0:
                new_width = max(100, int(width * resize_factor))
                new_height = max(100, int(height * resize_factor))
                image = image.resize((new_width, new_height), Image.LANCZOS)
            
            # 如果有目标尺寸，调整到目标尺寸
            if target_size:
                image = self.resize_and_crop(image, target_size)
            
            # 确保格式简单
            if image.mode not in ['RGB', 'RGBA']:
                image = image.convert('RGB')
            
            return image
            
        except Exception as e:
            self.logger.error(f"保守加载也失败 {image_path}: {e}")
            # 返回一个默认图像
            return Image.new('RGB', target_size or (300, 300), (128, 128, 128))
    
    # === 图像处理功能 (来自demo_refactored_templates.py) ===
    
    @staticmethod
    def resize_image(image: Union[Image.Image, np.ndarray], target_width: int) -> Image.Image:
        """
        调整图像尺寸 - 与原版完全一致
        
        Args:
            image: 输入图像
            target_width: 目标宽度
            
        Returns:
            调整后的图像
        """
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        if not isinstance(image, Image.Image):
            raise ValueError("输入必须是PIL图像或numpy数组")
        
        original_width, original_height = image.size
        
        if original_width == target_width:
            return image
        
        # 保持宽高比
        aspect_ratio = original_height / original_width
        target_height = int(target_width * aspect_ratio)
        
        # 使用高质量重采样
        resized_image = image.resize((target_width, target_height), Image.LANCZOS)
        
        return resized_image
    
    @staticmethod
    def ensure_rgba(image: Image.Image) -> Image.Image:
        """
        确保图像为RGBA格式 - 与原版完全一致
        
        Args:
            image: 输入图像
            
        Returns:
            RGBA格式的图像
        """
        if image.mode == 'RGBA':
            return image
        elif image.mode == 'RGB':
            return image.convert('RGBA')
        elif image.mode == 'P':
            # 调色板模式，检查是否有透明度
            if 'transparency' in image.info:
                return image.convert('RGBA')
            else:
                return image.convert('RGB').convert('RGBA')
        else:
            # 其他模式，先转RGB再转RGBA
            return image.convert('RGB').convert('RGBA')
    
    @staticmethod
    def resize_and_crop(image: Image.Image, target_size: Tuple[int, int]) -> Image.Image:
        """
        调整尺寸并裁剪 - 与原版算法完全一致
        
        Args:
            image: 输入图像
            target_size: 目标尺寸 (宽, 高)
            
        Returns:
            处理后的图像
        """
        target_width, target_height = target_size
        original_width, original_height = image.size
        
        # 计算缩放比例，选择较大的比例以确保覆盖目标尺寸
        scale_w = target_width / original_width
        scale_h = target_height / original_height
        scale = max(scale_w, scale_h)
        
        # 计算缩放后的尺寸
        new_width = int(original_width * scale)
        new_height = int(original_height * scale)
        
        # 缩放图像
        scaled_image = image.resize((new_width, new_height), Image.LANCZOS)
        
        # 计算裁剪区域 (居中裁剪)
        left = (new_width - target_width) // 2
        top = (new_height - target_height) // 2
        right = left + target_width
        bottom = top + target_height
        
        # 裁剪图像
        cropped_image = scaled_image.crop((left, top, right, bottom))
        
        return cropped_image
    
    @staticmethod
    def enhance_image(image: Union[Image.Image, np.ndarray], 
                     contrast: float = 1.0, 
                     brightness: float = 1.0, 
                     saturation: float = 1.0) -> Image.Image:
        """
        图像增强 - 与原版效果完全一致
        
        Args:
            image: 输入图像
            contrast: 对比度 (1.0为原始)
            brightness: 亮度 (1.0为原始)
            saturation: 饱和度 (1.0为原始)
            
        Returns:
            增强后的图像
        """
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        if not isinstance(image, Image.Image):
            raise ValueError("输入必须是PIL图像或numpy数组")
        
        # 确保图像格式
        if image.mode not in ['RGB', 'RGBA']:
            image = image.convert('RGB')
        
        enhanced = image
        
        # 调整亮度
        if brightness != 1.0:
            enhancer = ImageEnhance.Brightness(enhanced)
            enhanced = enhancer.enhance(brightness)
        
        # 调整对比度
        if contrast != 1.0:
            enhancer = ImageEnhance.Contrast(enhanced)
            enhanced = enhancer.enhance(contrast)
        
        # 调整饱和度
        if saturation != 1.0:
            enhancer = ImageEnhance.Color(enhanced)
            enhanced = enhancer.enhance(saturation)
        
        return enhanced
    
    @staticmethod
    def load_image(image_path: str) -> Optional[Image.Image]:
        """
        加载图像 - 与原版完全一致
        
        Args:
            image_path: 图像路径
            
        Returns:
            加载的图像或None
        """
        try:
            if not os.path.exists(image_path):
                logger.warning(f"图像文件不存在: {image_path}")
                return None
            
            image = Image.open(image_path)
            
            # 确保图像格式
            if image.mode not in ['RGB', 'RGBA']:
                if image.mode == 'P' and 'transparency' in image.info:
                    image = image.convert('RGBA')
                else:
                    image = image.convert('RGB')
            
            return image
            
        except Exception as e:
            logger.error(f"加载图像失败 {image_path}: {e}")
            return None
    
    @staticmethod
    def apply_blur(image: Image.Image, radius: float = 2.0) -> Image.Image:
        """
        应用模糊效果 - 与原版效果完全一致
        
        Args:
            image: 输入图像
            radius: 模糊半径
            
        Returns:
            模糊后的图像
        """
        return image.filter(ImageFilter.GaussianBlur(radius=radius))
    
    @staticmethod
    def create_canvas(size: Tuple[int, int], color: Tuple[int, int, int, int] = (255, 255, 255, 255)) -> Image.Image:
        """
        创建画布 - 与原版完全一致
        
        Args:
            size: 画布尺寸
            color: 背景颜色 (RGBA)
            
        Returns:
            新画布
        """
        return Image.new('RGBA', size, color)
    
    @staticmethod
    def paste_image(canvas: Image.Image, image: Image.Image, position: Tuple[int, int], mask: Optional[Image.Image] = None) -> Image.Image:
        """
        粘贴图像到画布 - 与原版完全一致
        
        Args:
            canvas: 目标画布
            image: 要粘贴的图像
            position: 粘贴位置 (x, y)
            mask: 可选的蒙版
            
        Returns:
            更新后的画布
        """
        # 确保画布是RGBA格式
        if canvas.mode != 'RGBA':
            canvas = canvas.convert('RGBA')
        
        # 确保图像兼容
        if image.mode not in ['RGB', 'RGBA']:
            image = image.convert('RGBA')
        
        # 粘贴图像
        if mask:
            canvas.paste(image, position, mask)
        elif image.mode == 'RGBA':
            canvas.paste(image, position, image)
        else:
            canvas.paste(image, position)
        
        return canvas
    
    @staticmethod
    def alpha_composite(base: Image.Image, overlay: Image.Image) -> Image.Image:
        """
        Alpha合成 - 与原版完全一致
        
        Args:
            base: 基础图像
            overlay: 覆盖图像
            
        Returns:
            合成后的图像
        """
        # 确保两个图像都是RGBA格式且尺寸相同
        if base.mode != 'RGBA':
            base = base.convert('RGBA')
        if overlay.mode != 'RGBA':
            overlay = overlay.convert('RGBA')
        
        # 调整覆盖图像尺寸以匹配基础图像
        if overlay.size != base.size:
            overlay = overlay.resize(base.size, Image.LANCZOS)
        
        # 执行Alpha合成
        return Image.alpha_composite(base, overlay)
    
    # === 批量处理功能 (来自poster_template.py) ===
    
    def batch_process_directory(self, directory: str, max_images: int = 100) -> list:
        """
        批量处理目录中的图像 - 与原版完全一致
        
        Args:
            directory: 图像目录
            max_images: 最大处理数量
            
        Returns:
            处理结果列表
        """
        if not os.path.exists(directory):
            raise FileNotFoundError(f"目录不存在: {directory}")
        
        # 支持的图像格式
        supported_formats = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
        
        results = []
        processed_count = 0
        
        try:
            # 获取所有图像文件
            image_files = []
            for file in os.listdir(directory):
                file_path = os.path.join(directory, file)
                if os.path.isfile(file_path):
                    ext = os.path.splitext(file)[1].lower()
                    if ext in supported_formats:
                        image_files.append(file_path)
            
            self.logger.info(f"找到 {len(image_files)} 个图像文件")
            
            # 限制处理数量
            if len(image_files) > max_images:
                image_files = image_files[:max_images]
                self.logger.info(f"限制处理数量为 {max_images}")
            
            # 批量处理
            for i, image_path in enumerate(image_files):
                # 检查内存压力
                if i % 10 == 0 and self.check_memory_pressure():
                    self.logger.warning("内存压力过高，清理缓存")
                    self.release_loaded_images()
                    self.force_garbage_collection()
                
                try:
                    # 处理单个图像
                    image = self.smart_load_image(image_path)
                    
                    result = {
                        'file_path': image_path,
                        'original_size': image.size,
                        'status': 'success',
                        'processed_at': i + 1
                    }
                    
                    results.append(result)
                    processed_count += 1
                    
                    if processed_count % 20 == 0:
                        memory_info = self.get_memory_usage()
                        self.logger.info(f"已处理 {processed_count}/{len(image_files)} 图像，内存使用: {memory_info[0]:.1f}%")
                    
                except Exception as e:
                    self.logger.error(f"处理图像失败 {image_path}: {e}")
                    results.append({
                        'file_path': image_path,
                        'status': 'error',
                        'error': str(e),
                        'processed_at': i + 1
                    })
            
            self.logger.info(f"批量处理完成 - 成功: {processed_count}, 失败: {len(results) - processed_count}")
            
        except Exception as e:
            self.logger.error(f"批量处理失败: {e}")
            raise
        
        finally:
            # 清理资源
            self.release_loaded_images()
            self.force_garbage_collection()
        
        return results
    
    def save_optimized_image(self, image: Image.Image, output_path: str, quality: int = 85) -> bool:
        """
        保存优化图像 - 与原版完全一致
        
        Args:
            image: 要保存的图像
            output_path: 输出路径
            quality: JPEG质量 (1-100)
            
        Returns:
            是否保存成功
        """
        try:
            # 创建输出目录
            output_dir = os.path.dirname(output_path)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir, exist_ok=True)
            
            # 根据文件扩展名确定格式
            ext = os.path.splitext(output_path)[1].lower()
            
            if ext in ['.jpg', '.jpeg']:
                # JPEG格式，需要转换为RGB
                if image.mode in ['RGBA', 'P']:
                    # 创建白色背景
                    rgb_image = Image.new('RGB', image.size, (255, 255, 255))
                    if image.mode == 'RGBA':
                        rgb_image.paste(image, mask=image.split()[-1])  # 使用Alpha通道作为蒙版
                    else:
                        rgb_image.paste(image)
                    image = rgb_image
                
                image.save(output_path, 'JPEG', quality=quality, optimize=True)
                
            elif ext == '.png':
                # PNG格式，保持透明度
                if image.mode not in ['RGB', 'RGBA']:
                    image = image.convert('RGBA')
                image.save(output_path, 'PNG', optimize=True)
                
            elif ext == '.webp':
                # WebP格式
                image.save(output_path, 'WEBP', quality=quality, method=6)
                
            else:
                # 其他格式，使用默认
                image.save(output_path)
            
            self.logger.debug(f"图像已保存: {output_path}")
            return True
            
        except Exception as e:
            self.logger.error(f"保存图像失败 {output_path}: {e}")
            return False 