import re
from konabot.plugins.fx_process.color_handle import ColorHandle
import numpy as np
from PIL import Image, ImageDraw
from typing import List, Tuple, Dict, Optional

class GradientGenerator:
    """渐变生成器类"""
    
    def __init__(self):
        self.has_numpy = hasattr(np, '__version__')
    
    def parse_color_list(self, color_list_str: str) -> List[Dict]:
        """解析渐变颜色列表字符串
        
        Args:
            color_list_str: 格式如 '[rgb(255,0,0)|(0,0)+rgb(0,255,0)|(0,100)+rgb(0,0,255)|(50,100)]'
        
        Returns:
            list: 包含颜色和位置信息的字典列表
        """
        color_nodes = []
        color_list_str = color_list_str.strip('[]').strip()
        matches = color_list_str.split('+')
        
        for single_str in matches:
            color_str = single_str.split('|')[0]
            pos_str = single_str.split('|')[1] if '|' in single_str else '0,0'

            color = ColorHandle.parse_color(color_str.strip())

            try:
                pos_str = pos_str.replace('(', '').replace(')', '')
                x_str, y_str = pos_str.split(',')
                x_percent = float(x_str.strip().replace('%', ''))
                y_percent = float(y_str.strip().replace('%', ''))
                x_percent = max(0, min(100, x_percent))
                y_percent = max(0, min(100, y_percent))
            except:
                x_percent = 0
                y_percent = 0
            
            color_nodes.append({
                'color': color,
                'position': (x_percent / 100.0, y_percent / 100.0)
            })
        
        if not color_nodes:
            color_nodes = [
                {'color': (255, 0, 0), 'position': (0, 0)},
                {'color': (0, 0, 255), 'position': (1, 1)}
            ]
        
        return color_nodes
    
    def create_gradient(self, width: int, height: int, color_nodes: List[Dict]) -> Image.Image:
        """创建渐变图像
        
        Args:
            width: 图像宽度
            height: 图像高度
            color_nodes: 颜色节点列表
        
        Returns:
            Image.Image: 渐变图像
        """
        if len(color_nodes) == 1:
            return Image.new('RGB', (width, height), color_nodes[0]['color'])
        elif len(color_nodes) == 2:
            return self._create_linear_gradient(width, height, color_nodes)
        else:
            return self._create_radial_gradient(width, height, color_nodes)
    
    def _create_linear_gradient(self, width: int, height: int, color_nodes: List[Dict]) -> Image.Image:
        """创建线性渐变"""
        color1 = color_nodes[0]['color']
        color2 = color_nodes[1]['color']
        pos1 = color_nodes[0]['position']
        pos2 = color_nodes[1]['position']
        
        if self.has_numpy:
            return self._create_linear_gradient_numpy(width, height, color1, color2, pos1, pos2)
        else:
            return self._create_linear_gradient_pil(width, height, color1, color2, pos1, pos2)
    
    def _create_linear_gradient_numpy(self, width: int, height: int, 
                                    color1: Tuple, color2: Tuple, 
                                    pos1: Tuple, pos2: Tuple) -> Image.Image:
        """使用numpy创建线性渐变"""
        # 创建坐标网格
        x = np.linspace(0, 1, width)
        y = np.linspace(0, 1, height)
        xx, yy = np.meshgrid(x, y)
        
        # 计算渐变方向
        dx = pos2[0] - pos1[0]
        dy = pos2[1] - pos1[1]
        length_sq = dx * dx + dy * dy
        
        if length_sq > 0:
            # 计算投影参数
            t = ((xx - pos1[0]) * dx + (yy - pos1[1]) * dy) / length_sq
            t = np.clip(t, 0, 1)
        else:
            t = np.zeros_like(xx)
        
        # 插值颜色
        r = color1[0] + (color2[0] - color1[0]) * t
        g = color1[1] + (color2[1] - color1[1]) * t
        b = color1[2] + (color2[2] - color1[2]) * t
        
        # 创建图像
        gradient_array = np.stack([r, g, b], axis=-1).astype(np.uint8)
        return Image.fromarray(gradient_array)
    
    def _create_linear_gradient_pil(self, width: int, height: int, 
                                   color1: Tuple, color2: Tuple, 
                                   pos1: Tuple, pos2: Tuple) -> Image.Image:
        """使用PIL创建线性渐变（没有numpy时使用）"""
        gradient = Image.new('RGB', (width, height))
        draw = ImageDraw.Draw(gradient)
        
        # 判断渐变方向
        if abs(pos1[0] - pos2[0]) < 0.01:  # 垂直渐变
            y1 = int(pos1[1] * (height - 1))
            y2 = int(pos2[1] * (height - 1))
            
            if y2 < y1:
                y1, y2 = y2, y1
                color1, color2 = color2, color1
            
            if y2 > y1:
                for y in range(height):
                    if y <= y1:
                        fill_color = color1
                    elif y >= y2:
                        fill_color = color2
                    else:
                        ratio = (y - y1) / (y2 - y1)
                        r = int(color1[0] + (color2[0] - color1[0]) * ratio)
                        g = int(color1[1] + (color2[1] - color1[1]) * ratio)
                        b = int(color1[2] + (color2[2] - color1[2]) * ratio)
                        fill_color = (r, g, b)
                    
                    draw.line([(0, y), (width, y)], fill=fill_color)
            else:
                draw.rectangle([0, 0, width, height], fill=color1)
                
        elif abs(pos1[1] - pos2[1]) < 0.01:  # 水平渐变
            x1 = int(pos1[0] * (width - 1))
            x2 = int(pos2[0] * (width - 1))
            
            if x2 < x1:
                x1, x2 = x2, x1
                color1, color2 = color2, color1
            
            if x2 > x1:
                for x in range(width):
                    if x <= x1:
                        fill_color = color1
                    elif x >= x2:
                        fill_color = color2
                    else:
                        ratio = (x - x1) / (x2 - x1)
                        r = int(color1[0] + (color2[0] - color1[0]) * ratio)
                        g = int(color1[1] + (color2[1] - color1[1]) * ratio)
                        b = int(color1[2] + (color2[2] - color1[2]) * ratio)
                        fill_color = (r, g, b)
                    
                    draw.line([(x, 0), (x, height)], fill=fill_color)
            else:
                draw.rectangle([0, 0, width, height], fill=color1)
                
        else:  # 对角渐变（简化处理为左上到右下）
            for y in range(height):
                for x in range(width):
                    distance = (x/width + y/height) / 2
                    r = int(color1[0] + (color2[0] - color1[0]) * distance)
                    g = int(color1[1] + (color2[1] - color1[1]) * distance)
                    b = int(color1[2] + (color2[2] - color1[2]) * distance)
                    draw.point((x, y), fill=(r, g, b))
        
        return gradient
    
    def _create_radial_gradient(self, width: int, height: int, color_nodes: List[Dict]) -> Image.Image:
        """创建径向渐变"""
        if self.has_numpy and len(color_nodes) > 2:
            return self._create_radial_gradient_numpy(width, height, color_nodes)
        else:
            return self._create_simple_gradient(width, height, color_nodes)
    
    def _create_radial_gradient_numpy(self, width: int, height: int, color_nodes: List[Dict]) -> Image.Image:
        """使用numpy创建径向渐变（多色）"""
        # 创建坐标网格
        x = np.linspace(0, 1, width)
        y = np.linspace(0, 1, height)
        xx, yy = np.meshgrid(x, y)
        
        # 提取颜色和位置
        positions = np.array([node['position'] for node in color_nodes])
        colors = np.array([node['color'] for node in color_nodes])
        
        # 计算每个点到所有节点的距离
        distances = np.sqrt((xx[:, :, np.newaxis] - positions[np.newaxis, np.newaxis, :, 0]) ** 2 +
                           (yy[:, :, np.newaxis] - positions[np.newaxis, np.newaxis, :, 1]) ** 2)
        
        # 找到最近的两个节点
        sorted_indices = np.argsort(distances, axis=2)
        nearest_idx = sorted_indices[:, :, 0]
        second_idx = sorted_indices[:, :, 1]
        
        # 获取对应的颜色
        nearest_colors = colors[nearest_idx]
        second_colors = colors[second_idx]
        
        # 获取距离并计算权重
        nearest_dist = np.take_along_axis(distances, np.expand_dims(nearest_idx, axis=2), axis=2)[:, :, 0]
        second_dist = np.take_along_axis(distances, np.expand_dims(second_idx, axis=2), axis=2)[:, :, 0]
        
        total_dist = nearest_dist + second_dist
        mask = total_dist > 0
        weight1 = np.zeros_like(nearest_dist)
        weight1[mask] = second_dist[mask] / total_dist[mask]
        weight2 = 1 - weight1
        
        # 插值颜色
        r = nearest_colors[:, :, 0] * weight1 + second_colors[:, :, 0] * weight2
        g = nearest_colors[:, :, 1] * weight1 + second_colors[:, :, 1] * weight2
        b = nearest_colors[:, :, 2] * weight1 + second_colors[:, :, 2] * weight2
        
        gradient_array = np.stack([r, g, b], axis=-1).astype(np.uint8)
        return Image.fromarray(gradient_array)
    
    def _create_simple_gradient(self, width: int, height: int, color_nodes: List[Dict]) -> Image.Image:
        """创建简化渐变（没有numpy或多色时使用）"""
        gradient = Image.new('RGB', (width, height))
        draw = ImageDraw.Draw(gradient)
        
        if len(color_nodes) >= 2:
            # 使用第一个和最后一个颜色创建简单渐变
            color1 = color_nodes[0]['color']
            color2 = color_nodes[-1]['color']
            
            # 判断节点分布
            x_positions = [node['position'][0] for node in color_nodes]
            y_positions = [node['position'][1] for node in color_nodes]
            
            if all(abs(x - x_positions[0]) < 0.01 for x in x_positions):
                # 垂直渐变
                for y in range(height):
                    ratio = y / (height - 1) if height > 1 else 0
                    r = int(color1[0] + (color2[0] - color1[0]) * ratio)
                    g = int(color1[1] + (color2[1] - color1[1]) * ratio)
                    b = int(color1[2] + (color2[2] - color1[2]) * ratio)
                    draw.line([(0, y), (width, y)], fill=(r, g, b))
            else:
                # 水平渐变
                for x in range(width):
                    ratio = x / (width - 1) if width > 1 else 0
                    r = int(color1[0] + (color2[0] - color1[0]) * ratio)
                    g = int(color1[1] + (color2[1] - color1[1]) * ratio)
                    b = int(color1[2] + (color2[2] - color1[2]) * ratio)
                    draw.line([(x, 0), (x, height)], fill=(r, g, b))
        else:
            # 单色
            draw.rectangle([0, 0, width, height], fill=color_nodes[0]['color'])
        
        return gradient
    
    def create_simple_gradient(self, width: int, height: int, 
                              start_color: Tuple, end_color: Tuple, 
                              direction: str = 'vertical') -> Image.Image:
        """创建简单双色渐变
        
        Args:
            width: 图像宽度
            height: 图像高度
            start_color: 起始颜色
            end_color: 结束颜色
            direction: 渐变方向 'vertical', 'horizontal', 'diagonal'
        
        Returns:
            Image.Image: 渐变图像
        """
        if direction == 'vertical':
            return self._create_vertical_gradient(width, height, start_color, end_color)
        elif direction == 'horizontal':
            return self._create_horizontal_gradient(width, height, start_color, end_color)
        else:  # diagonal
            return self._create_diagonal_gradient(width, height, start_color, end_color)
    
    def _create_vertical_gradient(self, width: int, height: int, 
                                 color1: Tuple, color2: Tuple) -> Image.Image:
        """创建垂直渐变"""
        gradient = Image.new('RGB', (width, height))
        draw = ImageDraw.Draw(gradient)
        
        for y in range(height):
            ratio = y / (height - 1) if height > 1 else 0
            r = int(color1[0] + (color2[0] - color1[0]) * ratio)
            g = int(color1[1] + (color2[1] - color1[1]) * ratio)
            b = int(color1[2] + (color2[2] - color1[2]) * ratio)
            draw.line([(0, y), (width, y)], fill=(r, g, b))
        
        return gradient
    
    def _create_horizontal_gradient(self, width: int, height: int, 
                                   color1: Tuple, color2: Tuple) -> Image.Image:
        """创建水平渐变"""
        gradient = Image.new('RGB', (width, height))
        draw = ImageDraw.Draw(gradient)
        
        for x in range(width):
            ratio = x / (width - 1) if width > 1 else 0
            r = int(color1[0] + (color2[0] - color1[0]) * ratio)
            g = int(color1[1] + (color2[1] - color1[1]) * ratio)
            b = int(color1[2] + (color2[2] - color1[2]) * ratio)
            draw.line([(x, 0), (x, height)], fill=(r, g, b))
        
        return gradient
    
    def _create_diagonal_gradient(self, width: int, height: int, 
                                 color1: Tuple, color2: Tuple) -> Image.Image:
        """创建对角渐变"""
        if self.has_numpy:
            return self._create_diagonal_gradient_numpy(width, height, color1, color2)
        else:
            return self._create_horizontal_gradient(width, height, color1, color2)  # 降级为水平渐变
    
    def _create_diagonal_gradient_numpy(self, width: int, height: int, 
                                       color1: Tuple, color2: Tuple) -> Image.Image:
        """使用numpy创建对角渐变"""
        x = np.linspace(0, 1, width)
        y = np.linspace(0, 1, height)
        xx, yy = np.meshgrid(x, y)
        
        distance = (xx + yy) / 2.0
        
        r = color1[0] + (color2[0] - color1[0]) * distance
        g = color1[1] + (color2[1] - color1[1]) * distance
        b = color1[2] + (color2[2] - color1[2]) * distance
        
        gradient_array = np.stack([r, g, b], axis=-1).astype(np.uint8)
        return Image.fromarray(gradient_array)