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

"""
颜色提取组件 - 新架构版本
完全保持与原版demo_refactored_templates.py中ColorExtractor的效果一致
"""

import logging
import colorsys
import math
from typing import Tuple, List, Dict, Any, Optional
from PIL import Image
import random
from .text_renderer import PosterComponent

logger = logging.getLogger(__name__)


class ColorExtractor(PosterComponent):
    """
    颜色提取器 - 完全保持原版效果
    提供智能颜色提取、配色方案生成、颜色调和等功能
    """
    
    # 预定义的颜色主题（顶部色，底部色）- 与原版完全一致
    COLOR_THEMES = {
        "ocean_deep": ((20, 51, 82), (13, 71, 161)),
        "sunset": ((255, 152, 0), (255, 87, 34)),
        "forest": ((56, 142, 60), (27, 94, 32)),
        "lavender": ((149, 117, 205), (106, 27, 154)),
        "coral": ((255, 138, 128), (244, 67, 54)),
        "sky": ((100, 181, 246), (33, 150, 243)),
        "earth": ((141, 110, 99), (93, 64, 55)),
        "golden": ((255, 193, 7), (255, 152, 0)),
        "mint": ((129, 199, 132), (76, 175, 80)),
        "rose": ((240, 98, 146), (233, 30, 99))
    }
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化颜色提取器"""
        # 可以从配置中添加自定义主题
        custom_themes = config.get('custom_themes', {})
        if custom_themes:
            self.COLOR_THEMES.update(custom_themes)
            
    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 == 'extract_dominant':
                return self.extract_dominant_color(**params)
            elif operation == 'generate_palette':
                return self.get_complementary_color(**params)
            elif operation == 'harmonize_colors':
                return self.ensure_colors_harmony(**params)
            elif operation == 'get_theme':
                return self.get_theme_colors(**params)
                
        return None
    
    def cleanup(self) -> None:
        """清理资源"""
        pass
    
    @staticmethod
    def extract_dominant_color(image: Image.Image, 
                             sample_size: int = 200, 
                             sample_method: str = "grid") -> Tuple[int, int, int]:
        """
        提取图像的主导颜色 - 与原版算法完全一致
        
        Args:
            image: PIL图像对象
            sample_size: 采样大小
            sample_method: 采样方法 ("grid" 或 "edge")
            
        Returns:
            RGB颜色元组
        """
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        width, height = image.size
        
        # 根据采样方法收集像素
        if sample_method == "edge":
            # 边缘采样 - 更适合提取背景色
            pixels = []
            
            # 采样边缘区域
            edge_width = min(width // 10, 50)  # 边缘宽度
            
            # 上边缘
            for y in range(min(edge_width, height)):
                for x in range(0, width, max(1, width // 20)):
                    pixels.append(image.getpixel((x, y)))
            
            # 下边缘
            for y in range(max(0, height - edge_width), height):
                for x in range(0, width, max(1, width // 20)):
                    pixels.append(image.getpixel((x, y)))
            
            # 左边缘
            for x in range(min(edge_width, width)):
                for y in range(0, height, max(1, height // 20)):
                    pixels.append(image.getpixel((x, y)))
            
            # 右边缘
            for x in range(max(0, width - edge_width), width):
                for y in range(0, height, max(1, height // 20)):
                    pixels.append(image.getpixel((x, y)))
                    
        else:
            # 网格采样 - 更均匀的采样
            step_x = max(1, width // int(math.sqrt(sample_size)))
            step_y = max(1, height // int(math.sqrt(sample_size)))
            
            pixels = []
            for y in range(0, height, step_y):
                for x in range(0, width, step_x):
                    pixels.append(image.getpixel((x, y)))
        
        if not pixels:
            return (128, 128, 128)  # 默认灰色
        
        # 使用改进的聚类算法找到主导颜色
        return ColorExtractor._find_dominant_color_improved(pixels)
    
    @staticmethod
    def _find_dominant_color_improved(pixels: List[Tuple[int, int, int]]) -> Tuple[int, int, int]:
        """
        改进的主导颜色查找算法 - 与原版完全一致
        """
        if not pixels:
            return (128, 128, 128)
        
        # 颜色直方图统计
        color_counts = {}
        
        for pixel in pixels:
            r, g, b = pixel
            
            # 过滤掉过于极端的颜色（纯黑、纯白、过饱和）
            if (r + g + b) < 30 or (r + g + b) > 735:  # 过暗或过亮
                continue
            if max(r, g, b) - min(r, g, b) > 200:  # 过饱和
                continue
            
            # 颜色量化以减少噪声
            quantized = (
                (r // 16) * 16,
                (g // 16) * 16, 
                (b // 16) * 16
            )
            
            color_counts[quantized] = color_counts.get(quantized, 0) + 1
        
        if not color_counts:
            # 如果过滤后没有合适的颜色，降低过滤标准
            for pixel in pixels:
                r, g, b = pixel
                quantized = (
                    (r // 32) * 32,
                    (g // 32) * 32,
                    (b // 32) * 32
                )
                color_counts[quantized] = color_counts.get(quantized, 0) + 1
        
        # 找到出现频率最高的颜色
        if color_counts:
            dominant_color = max(color_counts.items(), key=lambda x: x[1])[0]
            return dominant_color
        
        # 最后的备选方案：计算平均颜色
        avg_r = sum(p[0] for p in pixels) // len(pixels)
        avg_g = sum(p[1] for p in pixels) // len(pixels)
        avg_b = sum(p[2] for p in pixels) // len(pixels)
        
        return (avg_r, avg_g, avg_b)
    
    @staticmethod
    def select_best_color(color_candidates: List[Tuple[Tuple[int, int, int], int]]) -> Tuple[int, int, int]:
        """
        从候选颜色中选择最佳颜色 - 与原版逻辑完全一致
        
        Args:
            color_candidates: [(颜色, 权重), ...] 的列表
            
        Returns:
            最佳颜色
        """
        if not color_candidates:
            return (100, 149, 237)  # 默认蓝色
        
        # 按权重排序
        sorted_candidates = sorted(color_candidates, key=lambda x: x[1], reverse=True)
        
        # 选择权重最高且颜色合适的候选
        for color, weight in sorted_candidates:
            r, g, b = color
            
            # 颜色质量评估
            brightness = (r + g + b) / 3
            saturation = (max(r, g, b) - min(r, g, b)) / max(max(r, g, b), 1)
            
            # 偏好中等亮度、适中饱和度的颜色
            if 30 <= brightness <= 200 and 0.1 <= saturation <= 0.8:
                return color
        
        # 如果没有找到合适的，返回权重最高的
        return sorted_candidates[0][0]
    
    @staticmethod
    def adjust_color_for_background(color: Tuple[int, int, int]) -> Tuple[int, int, int]:
        """
        调整颜色使其适合作为背景 - 与原版逻辑完全一致
        
        Args:
            color: 原始RGB颜色
            
        Returns:
            调整后的RGB颜色
        """
        r, g, b = color
        
        # 转换到HSV空间进行调整
        h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
        
        # 降低饱和度，使颜色更柔和
        s = min(s, 0.7)  # 饱和度不超过70%
        
        # 调整亮度，确保在合适范围内
        if v < 0.3:
            v = 0.3  # 最低30%亮度
        elif v > 0.8:
            v = 0.8  # 最高80%亮度
        
        # 转换回RGB
        r_new, g_new, b_new = colorsys.hsv_to_rgb(h, s, v)
        
        return (
            int(r_new * 255),
            int(g_new * 255), 
            int(b_new * 255)
        )
    
    @staticmethod
    def get_default_color() -> Tuple[int, int, int]:
        """
        获取默认颜色 - 与原版完全一致
        
        Returns:
            默认的蓝色
        """
        return (100, 149, 237)  # Cornflower Blue
    
    @staticmethod
    def ensure_colors_harmony(top_color: Tuple[int, int, int], 
                            bottom_color: Tuple[int, int, int], 
                            harmony_threshold: int = 30) -> Tuple[Tuple[int, int, int], Tuple[int, int, int]]:
        """
        确保两个颜色和谐 - 与原版算法完全一致
        
        Args:
            top_color: 顶部颜色
            bottom_color: 底部颜色
            harmony_threshold: 和谐度阈值
            
        Returns:
            调和后的(顶部颜色, 底部颜色)
        """
        def color_distance(c1, c2):
            """计算颜色距离"""
            return math.sqrt(sum((a - b) ** 2 for a, b in zip(c1, c2)))
        
        def adjust_color_harmony(base_color, target_color, factor=0.3):
            """调整颜色使其更和谐"""
            return tuple(
                int(base_color[i] + (target_color[i] - base_color[i]) * factor)
                for i in range(3)
            )
        
        # 检查当前颜色是否已经和谐
        if color_distance(top_color, bottom_color) <= harmony_threshold:
            return top_color, bottom_color
        
        # 转换到HSV空间进行调整
        top_hsv = colorsys.rgb_to_hsv(top_color[0]/255.0, top_color[1]/255.0, top_color[2]/255.0)
        bottom_hsv = colorsys.rgb_to_hsv(bottom_color[0]/255.0, bottom_color[1]/255.0, bottom_color[2]/255.0)
        
        # 调整色相，使其更接近
        hue_diff = abs(top_hsv[0] - bottom_hsv[0])
        if hue_diff > 0.5:  # 如果色相差距过大
            hue_diff = 1.0 - hue_diff
        
        if hue_diff > 0.25:  # 如果色相差距超过90度
            # 让底部颜色向顶部颜色的色相靠近
            if top_hsv[0] > bottom_hsv[0]:
                if top_hsv[0] - bottom_hsv[0] > 0.5:
                    new_bottom_hue = bottom_hsv[0] - 0.1
                else:
                    new_bottom_hue = bottom_hsv[0] + 0.1
            else:
                if bottom_hsv[0] - top_hsv[0] > 0.5:
                    new_bottom_hue = bottom_hsv[0] + 0.1
                else:
                    new_bottom_hue = bottom_hsv[0] - 0.1
            
            # 确保色相在0-1范围内
            new_bottom_hue = new_bottom_hue % 1.0
            
            # 创建新的底部颜色
            new_bottom_rgb = colorsys.hsv_to_rgb(new_bottom_hue, bottom_hsv[1], bottom_hsv[2])
            bottom_color = (
                int(new_bottom_rgb[0] * 255),
                int(new_bottom_rgb[1] * 255),
                int(new_bottom_rgb[2] * 255)
            )
        
        return top_color, bottom_color
    
    @staticmethod
    def get_theme_colors(theme_name: Optional[str] = None) -> Tuple[Tuple[int, int, int], Tuple[int, int, int]]:
        """
        获取主题颜色 - 与原版完全一致
        
        Args:
            theme_name: 主题名称，如果为None则随机选择
            
        Returns:
            (顶部颜色, 底部颜色)
        """
        if theme_name and theme_name in ColorExtractor.COLOR_THEMES:
            return ColorExtractor.COLOR_THEMES[theme_name]
        
        # 随机选择一个主题
        theme_name = random.choice(list(ColorExtractor.COLOR_THEMES.keys()))
        return ColorExtractor.COLOR_THEMES[theme_name]
    
    @staticmethod
    def create_gradient_colors(base_color: Tuple[int, int, int], 
                             variation: float = 0.3) -> Tuple[Tuple[int, int, int], Tuple[int, int, int]]:
        """
        基于基础颜色创建渐变色 - 与原版算法完全一致
        
        Args:
            base_color: 基础颜色
            variation: 变化程度 (0.0-1.0)
            
        Returns:
            (顶部颜色, 底部颜色)
        """
        r, g, b = base_color
        
        # 转换到HSV空间
        h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
        
        # 创建更亮的顶部颜色
        top_v = min(1.0, v + variation * 0.5)
        top_s = max(0.0, s - variation * 0.2)
        
        # 创建更暗的底部颜色
        bottom_v = max(0.0, v - variation * 0.3)
        bottom_s = min(1.0, s + variation * 0.1)
        
        # 转换回RGB
        top_rgb = colorsys.hsv_to_rgb(h, top_s, top_v)
        bottom_rgb = colorsys.hsv_to_rgb(h, bottom_s, bottom_v)
        
        top_color = (int(top_rgb[0] * 255), int(top_rgb[1] * 255), int(top_rgb[2] * 255))
        bottom_color = (int(bottom_rgb[0] * 255), int(bottom_rgb[1] * 255), int(bottom_rgb[2] * 255))
        
        return top_color, bottom_color
    
    @staticmethod
    def get_complementary_color(color: Tuple[int, int, int]) -> Tuple[int, int, int]:
        """
        获取补色 - 与原版算法完全一致
        
        Args:
            color: 原始RGB颜色
            
        Returns:
            补色RGB
        """
        r, g, b = color
        
        # 转换到HSV空间
        h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
        
        # 计算补色（色相相差180度）
        complement_h = (h + 0.5) % 1.0
        
        # 转换回RGB
        complement_rgb = colorsys.hsv_to_rgb(complement_h, s, v)
        
        return (
            int(complement_rgb[0] * 255),
            int(complement_rgb[1] * 255),
            int(complement_rgb[2] * 255)
        )
    
    @staticmethod
    def get_analogous_colors(color: Tuple[int, int, int], count: int = 2) -> List[Tuple[int, int, int]]:
        """
        获取类似色 - 与原版算法完全一致
        
        Args:
            color: 基础颜色
            count: 类似色数量
            
        Returns:
            类似色列表
        """
        r, g, b = color
        h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
        
        analogous_colors = []
        
        # 在色相轮上取相邻的颜色
        hue_step = 0.08  # 约30度
        
        for i in range(1, count + 1):
            # 一半向左，一半向右
            if i % 2 == 1:
                new_h = (h + (i // 2 + 1) * hue_step) % 1.0
            else:
                new_h = (h - (i // 2) * hue_step) % 1.0
            
            # 稍微调整饱和度和亮度以增加变化
            new_s = max(0.0, min(1.0, s + (random.random() - 0.5) * 0.1))
            new_v = max(0.0, min(1.0, v + (random.random() - 0.5) * 0.1))
            
            # 转换回RGB
            new_rgb = colorsys.hsv_to_rgb(new_h, new_s, new_v)
            analogous_color = (
                int(new_rgb[0] * 255),
                int(new_rgb[1] * 255),
                int(new_rgb[2] * 255)
            )
            analogous_colors.append(analogous_color)
        
        return analogous_colors
    
    @staticmethod
    def get_smart_feature_color(background_colors: Tuple[Tuple[int, int, int], Tuple[int, int, int]]) -> Tuple[int, int, int]:
        """
        根据背景颜色智能选择特征色 - 与原版算法完全一致
        
        Args:
            background_colors: (顶部颜色, 底部颜色)
            
        Returns:
            智能特征色
        """
        top_color, bottom_color = background_colors
        
        # 计算背景的平均颜色
        avg_r = (top_color[0] + bottom_color[0]) // 2
        avg_g = (top_color[1] + bottom_color[1]) // 2
        avg_b = (top_color[2] + bottom_color[2]) // 2
        avg_color = (avg_r, avg_g, avg_b)
        
        # 转换到HSV
        h, s, v = colorsys.rgb_to_hsv(avg_r/255.0, avg_g/255.0, avg_b/255.0)
        
        def get_color_tone(color):
            """判断颜色的色调倾向"""
            r, g, b = color
            if r > g and r > b:
                return "warm"  # 暖色调
            elif b > r and b > g:
                return "cool"  # 冷色调
            else:
                return "neutral"  # 中性色调
        
        # 根据背景色调选择对比色
        bg_tone = get_color_tone(avg_color)
        
        if bg_tone == "warm":
            # 暖色背景选择冷色特征
            feature_h = (h + 0.3) % 1.0  # 偏向蓝绿色
        elif bg_tone == "cool":
            # 冷色背景选择暖色特征  
            feature_h = (h + 0.7) % 1.0  # 偏向橙红色
        else:
            # 中性背景选择互补色
            feature_h = (h + 0.5) % 1.0
        
        # 调整饱和度和亮度以确保对比度
        feature_s = max(0.6, min(1.0, 1.0 - s * 0.5))  # 高饱和度
        feature_v = 0.8 if v < 0.5 else 0.3  # 与背景亮度相反
        
        # 转换回RGB
        feature_rgb = colorsys.hsv_to_rgb(feature_h, feature_s, feature_v)
        
        return (
            int(feature_rgb[0] * 255),
            int(feature_rgb[1] * 255),
            int(feature_rgb[2] * 255)
        ) 