"""
ComfyUI节点实现
"""
import torch
import numpy as np
import time
import os
from PIL import Image
from .api_client import GQCHAPIClient


def tensor_to_pil(tensor):
    """将ComfyUI的tensor转换为PIL Image"""
    # tensor shape: [B, H, W, C]
    if len(tensor.shape) == 4:
        tensor = tensor[0]  # 取第一张图片
    
    # 转换为numpy数组
    np_image = tensor.cpu().numpy()
    
    # 确保值在0-255范围内
    if np_image.dtype == np.float32 or np_image.dtype == np.float64:
        np_image = (np_image * 255).astype(np.uint8)
    
    # 转换为PIL Image
    image = Image.fromarray(np_image)
    return image


def pil_to_tensor(image):
    """将PIL Image转换为ComfyUI的tensor"""
    # 转换为numpy数组
    np_image = np.array(image).astype(np.float32) / 255.0
    
    # 转换为tensor，shape: [1, H, W, C]
    tensor = torch.from_numpy(np_image).unsqueeze(0)
    return tensor


def wait_for_task_completion(client, task_id, check_interval=3, max_wait_time=600):
    """
    等待任务完成
    
    Args:
        client: API客户端
        task_id: 任务ID
        check_interval: 检查间隔（秒）
        max_wait_time: 最大等待时间（秒，默认600秒）
        
    Returns:
        任务结果字典
    """
    start_time = time.time()
    attempt = 0
    
    while True:
        attempt += 1
        elapsed_time = time.time() - start_time
        
        # 检查是否超时
        if elapsed_time > max_wait_time:
            raise Exception(f"任务超时：等待时间超过{max_wait_time}秒")
        
        # 查询任务状态
        try:
            result = client.check_task_status(task_id)
            status = result.get('status', 'Unknown')
            
            print(f"⏳ 第{attempt}次检查 (已等待{int(elapsed_time)}秒) - 状态: {status}")
            
            if status == 'Completed':
                print(f"✅ 任务完成！总耗时: {int(elapsed_time)}秒")
                return result
            elif status == 'Failed':
                raise Exception(f"任务失败: {result.get('error', '未知错误')}")
            elif status in [None, 'Error']:
                raise Exception(f"任务状态异常: {status}")
            
            # 等待后继续检查
            time.sleep(check_interval)
            
        except Exception as e:
            if "任务失败" in str(e) or "任务状态异常" in str(e) or "任务超时" in str(e):
                raise e
            else:
                print(f"⚠️  查询状态时出错: {str(e)}，{check_interval}秒后重试...")
                time.sleep(check_interval)


def download_vector_file(client, vector_url, task_id):
    """
    下载矢量图文件并保存到本地
    
    Args:
        client: API客户端
        vector_url: 矢量图URL
        task_id: 任务ID
        
    Returns:
        本地文件路径，如果没有矢量图则返回空字符串
    """
    if not vector_url:
        return ""
    
    try:
        # 创建输出目录
        output_dir = os.path.join(os.path.dirname(__file__), "output", "vectors")
        os.makedirs(output_dir, exist_ok=True)
        
        # 生成文件名
        filename = f"vector_{task_id}.eps"
        filepath = os.path.join(output_dir, filename)
        
        # 下载文件
        print(f"📥 下载矢量图文件...")
        response = client.session.get(vector_url, timeout=30)
        response.raise_for_status()
        
        # 保存文件
        with open(filepath, 'wb') as f:
            f.write(response.content)
        
        print(f"✅ 矢量图已保存到: {filepath}")
        return filepath
        
    except Exception as e:
        print(f"⚠️  下载矢量图失败: {str(e)}")
        return ""


class GQCHHDRedraw:
    """高清重绘任务节点（一体化，自动等待完成）"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "sdxl_model": (["通用模型", "真实模型"],),
                "style": ([
                    "通用风格",
                    "花卉风格",
                    "水彩风格",
                    "3D立体风格",
                    "绣花刺绣",
                    "牛仔纹路",
                    "卡通人物",
                    "油画风格",
                    "豹纹风格",
                    "抽象艺术",
                    "几何图案",
                    "皮带链条",
                    "条纹风格",
                    "风景照片",
                    "矢量图风格",
                ],),
                "upscale": ("FLOAT", {
                    "default": 1.0,
                    "min": 1.0,
                    "max": 8.0,
                    "step": 0.5
                }),
                "s_noise": ("FLOAT", {
                    "default": 1.003,
                    "min": 1.0,
                    "max": 1.1,
                    "step": 0.001
                }),
                "steps": ("INT", {
                    "default": 40,
                    "min": 20,
                    "max": 40,
                    "step": 1
                }),
                 "auto_generate_prompt": ("BOOLEAN", {
                    "default": False
                }),
                "prompt": ("STRING", {
                    "default": "",
                    "multiline": True
                }),
                "neg_prompt": ("STRING", {
                    "default": "",
                    "multiline": True
                }),
            },
            "optional": {
                "high_fidelity": ("BOOLEAN", {
                    "default": False
                }),
                
                "strong_inpaint": ("BOOLEAN", {
                    "default": False
                }),
               
                "strong_inpaint_intensity": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.1
                }),
                "vector_switch": ("BOOLEAN", {
                    "default": False
                }),
                "color_count": ("INT", {
                    "default": 5,
                    "min": 2,
                    "max": 16,
                    "step": 1
                }),
                
                "user_key": ("STRING", {
                    "default": ""
                }),
                "max_wait_time": ("INT", {
                    "default": 600,
                    "min": 60,
                    "max": 1800,
                    "step": 30
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE", "IMAGE", "STRING", "STRING")
    RETURN_NAMES = ("image", "image_1m", "vector_url", "info")
    FUNCTION = "process"
    CATEGORY = "GQCH Online"
    
    def process(self, image, sdxl_model, style, upscale, s_noise, steps, 
                auto_generate_prompt, prompt, neg_prompt,
                high_fidelity=False, strong_inpaint=False, strong_inpaint_intensity=0.5,
                vector_switch=False, color_count=5,
                user_key="", max_wait_time=600):
        """提交任务并等待完成"""
        try:
            client = GQCHAPIClient(user_key=user_key if user_key else None)
            pil_image = tensor_to_pil(image)
            
            print(f"📤 提交高清重绘任务...")
            result = client.submit_hd_redraw_task(
                image=pil_image,
                sdxl_model_dropdown=sdxl_model,
                style_dropdown=style,
                upscale=upscale,
                s_noise=s_noise,
                steps=steps,
                high_fidelity=high_fidelity,
                strong_inpaint=strong_inpaint,
                strong_inpaint_intensity=strong_inpaint_intensity,
                prompt=prompt,
                neg_prompt=neg_prompt,
                vector_switch=vector_switch,
                color_count=color_count,
                auto_generate_prompt=auto_generate_prompt
            )
            
            task_id = result['task_id']
            print(f"✅ 任务已提交，任务ID: {task_id}")
            print(f"📊 剩余次数: {result['remaining_count']}, 剩余金额: {result['remaining_money']}分")
            
            print(f"⏳ 等待任务完成（最长等待{max_wait_time}秒）...")
            completed_result = wait_for_task_completion(client, task_id, check_interval=3, max_wait_time=max_wait_time)
            
            image_url = completed_result.get('images_url')
            if not image_url:
                raise Exception("任务完成但没有返回图片URL")
            
            print(f"📥 下载处理后的图片...")
            output_image = client.download_image(image_url)
            output_tensor = pil_to_tensor(output_image)
            
            # 下载1倍图（如果有）
            image_url_1m = completed_result.get('images_url_1m', '')
            if image_url_1m:
                print(f"📥 下载1倍图...")
                output_image_1m = client.download_image(image_url_1m)
                output_tensor_1m = pil_to_tensor(output_image_1m)
            else:
                # 如果没有1倍图，返回原图
                output_tensor_1m = output_tensor
            
            # 获取矢量图URL（不下载）
            vector_url = completed_result.get('vector_url', '')
            if vector_url:
                print(f"📋 矢量图URL: {vector_url}")
            
            # 构建信息
            info_parts = [
                f"任务ID: {task_id}",
                f"剩余次数: {completed_result.get('remaining_count', 'N/A')}",
                f"剩余金额: {completed_result.get('remaining_money', 'N/A')}分"
            ]
            
            if completed_result.get('generated_prompt'):
                info_parts.append(f"生成的提示词: {completed_result.get('generated_prompt')}")
            
            info = ", ".join(info_parts)
            
            print(f"✅ 高清重绘任务完成！")
            return (output_tensor, output_tensor_1m, vector_url, info)
            
        except Exception as e:
            print(f"❌ 高清重绘任务失败: {str(e)}")
            raise e


class GQCHDownloadVector:
    """下载矢量图文件节点"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "vector_url": ("STRING", {
                    "default": "",
                    "multiline": False
                }),
            },
            "optional": {
                "filename": ("STRING", {
                    "default": "",
                    "multiline": False
                }),
            }
        }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("vector_path",)
    FUNCTION = "download"
    CATEGORY = "GQCH Online"
    
    def download(self, vector_url, filename=""):
        """下载矢量图文件"""
        try:
            if not vector_url:
                print("⚠️  没有提供矢量图URL")
                return ("",)
            
            # 创建输出目录
            output_dir = os.path.join(os.path.dirname(__file__), "output", "vectors")
            os.makedirs(output_dir, exist_ok=True)
            
            # 生成文件名
            if not filename:
                # 从URL中提取任务ID或使用时间戳
                import time
                timestamp = int(time.time())
                filename = f"vector_{timestamp}.eps"
            elif not filename.endswith('.eps'):
                filename = f"{filename}.eps"
            
            filepath = os.path.join(output_dir, filename)
            
            # 下载文件
            print(f"📥 下载矢量图文件...")
            import requests
            response = requests.get(vector_url, timeout=30)
            response.raise_for_status()
            
            # 保存文件
            with open(filepath, 'wb') as f:
                f.write(response.content)
            
            print(f"✅ 矢量图已保存到: {filepath}")
            return (filepath,)
            
        except Exception as e:
            print(f"❌ 下载矢量图失败: {str(e)}")
            return ("",)


class GQCHEditImage:
    """改图任务节点（一体化，自动等待完成）"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "prompt": ("STRING", {
                    "default": "",
                    "multiline": True
                }),
                "model_version": (["v1", "v2"],),
            },
            "optional": {
                "user_key": ("STRING", {
                    "default": ""
                }),
                "max_wait_time": ("INT", {
                    "default": 600,
                    "min": 60,
                    "max": 1800,
                    "step": 30
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("image", "info")
    FUNCTION = "process"
    CATEGORY = "GQCH Online"
    
    def process(self, image, prompt, model_version, user_key="", max_wait_time=600):
        """提交任务并等待完成"""
        try:
            # 初始化API客户端
            client = GQCHAPIClient(user_key=user_key if user_key else None)
            
            # 转换图片格式
            pil_image = tensor_to_pil(image)
            
            # 提交任务
            print(f"📤 提交改图任务...")
            result = client.submit_edit_image_task(
                image=pil_image,
                prompt=prompt,
                model_version=model_version
            )
            
            task_id = result['task_id']
            print(f"✅ 任务已提交，任务ID: {task_id}")
            print(f"📊 剩余次数: {result['remaining_count']}, 剩余金额: {result['remaining_money']}分")
            
            # 等待任务完成
            print(f"⏳ 等待任务完成（最长等待{max_wait_time}秒）...")
            completed_result = wait_for_task_completion(client, task_id, check_interval=3, max_wait_time=max_wait_time)
            
            # 获取图片URL
            image_url = completed_result.get('images_url')
            if not image_url:
                raise Exception("任务完成但没有返回图片URL")
            
            # 下载图片
            print(f"📥 下载处理后的图片...")
            output_image = client.download_image(image_url)
            
            # 转换为tensor
            output_tensor = pil_to_tensor(output_image)
            
            # 构建信息
            info = f"任务ID: {task_id}, 剩余次数: {completed_result.get('remaining_count', 'N/A')}, 剩余金额: {completed_result.get('remaining_money', 'N/A')}分"
            
            print(f"✅ 改图任务完成！")
            return (output_tensor, info)
            
        except Exception as e:
            print(f"❌ 改图任务失败: {str(e)}")
            raise e


class GQCHExpandImage:
    """扩图任务节点（一体化，自动等待完成）"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "expand_top": ("FLOAT", {
                    "default": 0.0,
                    "min": 0.0,
                    "max": 0.3,
                    "step": 0.05
                }),
                "expand_bottom": ("FLOAT", {
                    "default": 0.0,
                    "min": 0.0,
                    "max": 0.3,
                    "step": 0.05
                }),
                "expand_left": ("FLOAT", {
                    "default": 0.0,
                    "min": 0.0,
                    "max": 0.3,
                    "step": 0.05
                }),
                "expand_right": ("FLOAT", {
                    "default": 0.0,
                    "min": 0.0,
                    "max": 0.3,
                    "step": 0.05
                }),
                "prompt": ("STRING", {
                    "default": "",
                    "multiline": True
                }),
            },
            "optional": {
                "user_key": ("STRING", {
                    "default": ""
                }),
                "max_wait_time": ("INT", {
                    "default": 600,
                    "min": 60,
                    "max": 1800,
                    "step": 30
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("image", "info")
    FUNCTION = "process"
    CATEGORY = "GQCH Online"
    
    def process(self, image, expand_top, expand_bottom, expand_left, expand_right, prompt, user_key="", max_wait_time=600):
        """提交任务并等待完成"""
        try:
            client = GQCHAPIClient(user_key=user_key if user_key else None)
            pil_image = tensor_to_pil(image)
            
            print(f"📤 提交扩图任务...")
            result = client.submit_expand_image_task(
                image=pil_image,
                expand_top=expand_top,
                expand_bottom=expand_bottom,
                expand_left=expand_left,
                expand_right=expand_right,
                prompt=prompt
            )
            
            task_id = result['task_id']
            print(f"✅ 任务已提交，任务ID: {task_id}")
            print(f"📊 剩余次数: {result['remaining_count']}, 剩余金额: {result['remaining_money']}分")
            
            print(f"⏳ 等待任务完成（最长等待{max_wait_time}秒）...")
            completed_result = wait_for_task_completion(client, task_id, check_interval=3, max_wait_time=max_wait_time)
            
            image_url = completed_result.get('images_url')
            if not image_url:
                raise Exception("任务完成但没有返回图片URL")
            
            print(f"📥 下载处理后的图片...")
            output_image = client.download_image(image_url)
            output_tensor = pil_to_tensor(output_image)
            
            info = f"任务ID: {task_id}, 剩余次数: {completed_result.get('remaining_count', 'N/A')}, 剩余金额: {completed_result.get('remaining_money', 'N/A')}分"
            
            print(f"✅ 扩图任务完成！")
            return (output_tensor, info)
            
        except Exception as e:
            print(f"❌ 扩图任务失败: {str(e)}")
            raise e


class GQCHRemoveFabric:
    """移除布纹任务节点（一体化，自动等待完成）"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
            },
            "optional": {
                "user_key": ("STRING", {
                    "default": ""
                }),
                "max_wait_time": ("INT", {
                    "default": 600,
                    "min": 60,
                    "max": 1800,
                    "step": 30
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("image", "info")
    FUNCTION = "process"
    CATEGORY = "GQCH Online"
    
    def process(self, image, user_key="", max_wait_time=600):
        """提交任务并等待完成"""
        try:
            client = GQCHAPIClient(user_key=user_key if user_key else None)
            pil_image = tensor_to_pil(image)
            
            print(f"📤 提交移除布纹任务...")
            result = client.submit_remove_fabric_task(image=pil_image)
            
            task_id = result['task_id']
            print(f"✅ 任务已提交，任务ID: {task_id}")
            print(f"📊 剩余次数: {result['remaining_count']}, 剩余金额: {result['remaining_money']}分")
            
            print(f"⏳ 等待任务完成（最长等待{max_wait_time}秒）...")
            completed_result = wait_for_task_completion(client, task_id, check_interval=3, max_wait_time=max_wait_time)
            
            image_url = completed_result.get('images_url')
            if not image_url:
                raise Exception("任务完成但没有返回图片URL")
            
            print(f"📥 下载处理后的图片...")
            output_image = client.download_image(image_url)
            output_tensor = pil_to_tensor(output_image)
            
            info = f"任务ID: {task_id}, 剩余次数: {completed_result.get('remaining_count', 'N/A')}, 剩余金额: {completed_result.get('remaining_money', 'N/A')}分"
            
            print(f"✅ 移除布纹任务完成！")
            return (output_tensor, info)
            
        except Exception as e:
            print(f"❌ 移除布纹任务失败: {str(e)}")
            raise e


class GQCHRemoveDiamonds:
    """移除烫钻任务节点（一体化，自动等待完成）"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
            },
            "optional": {
                "user_key": ("STRING", {
                    "default": ""
                }),
                "max_wait_time": ("INT", {
                    "default": 600,
                    "min": 60,
                    "max": 1800,
                    "step": 30
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("image", "info")
    FUNCTION = "process"
    CATEGORY = "GQCH Online"
    
    def process(self, image, user_key="", max_wait_time=300):
        """提交任务并等待完成"""
        try:
            client = GQCHAPIClient(user_key=user_key if user_key else None)
            pil_image = tensor_to_pil(image)
            
            print(f"📤 提交移除烫钻任务...")
            result = client.submit_remove_diamonds_task(image=pil_image)
            
            task_id = result['task_id']
            print(f"✅ 任务已提交，任务ID: {task_id}")
            print(f"📊 剩余次数: {result['remaining_count']}, 剩余金额: {result['remaining_money']}分")
            
            print(f"⏳ 等待任务完成（最长等待{max_wait_time}秒）...")
            completed_result = wait_for_task_completion(client, task_id, check_interval=3, max_wait_time=max_wait_time)
            
            image_url = completed_result.get('images_url')
            if not image_url:
                raise Exception("任务完成但没有返回图片URL")
            
            print(f"📥 下载处理后的图片...")
            output_image = client.download_image(image_url)
            output_tensor = pil_to_tensor(output_image)
            
            info = f"任务ID: {task_id}, 剩余次数: {completed_result.get('remaining_count', 'N/A')}, 剩余金额: {completed_result.get('remaining_money', 'N/A')}分"
            
            print(f"✅ 移除烫钻任务完成！")
            return (output_tensor, info)
            
        except Exception as e:
            print(f"❌ 移除烫钻任务失败: {str(e)}")
            raise e


class GQCHRemoveWatermark:
    """去水印任务节点（一体化，自动等待完成）"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
            },
            "optional": {
                "user_key": ("STRING", {
                    "default": ""
                }),
                "max_wait_time": ("INT", {
                    "default": 600,
                    "min": 60,
                    "max": 1800,
                    "step": 30
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("image", "info")
    FUNCTION = "process"
    CATEGORY = "GQCH Online"
    
    def process(self, image, user_key="", max_wait_time=300):
        """提交任务并等待完成"""
        try:
            client = GQCHAPIClient(user_key=user_key if user_key else None)
            pil_image = tensor_to_pil(image)
            
            print(f"📤 提交去水印任务...")
            result = client.submit_remove_watermark_task(image=pil_image)
            
            task_id = result['task_id']
            print(f"✅ 任务已提交，任务ID: {task_id}")
            print(f"📊 剩余次数: {result['remaining_count']}, 剩余金额: {result['remaining_money']}分")
            
            print(f"⏳ 等待任务完成（最长等待{max_wait_time}秒）...")
            completed_result = wait_for_task_completion(client, task_id, check_interval=3, max_wait_time=max_wait_time)
            
            image_url = completed_result.get('images_url')
            if not image_url:
                raise Exception("任务完成但没有返回图片URL")
            
            print(f"📥 下载处理后的图片...")
            output_image = client.download_image(image_url)
            output_tensor = pil_to_tensor(output_image)
            
            info = f"任务ID: {task_id}, 剩余次数: {completed_result.get('remaining_count', 'N/A')}, 剩余金额: {completed_result.get('remaining_money', 'N/A')}分"
            
            print(f"✅ 去水印任务完成！")
            return (output_tensor, info)
            
        except Exception as e:
            print(f"❌ 去水印任务失败: {str(e)}")
            raise e


# 节点映射
NODE_CLASS_MAPPINGS = {
    "GQCHHDRedraw": GQCHHDRedraw,
    "GQCHDownloadVector": GQCHDownloadVector,
    "GQCHEditImage": GQCHEditImage,
    "GQCHExpandImage": GQCHExpandImage,
    "GQCHRemoveFabric": GQCHRemoveFabric,
    "GQCHRemoveDiamonds": GQCHRemoveDiamonds,
    "GQCHRemoveWatermark": GQCHRemoveWatermark,
}

# 节点显示名称
NODE_DISPLAY_NAME_MAPPINGS = {
    "GQCHHDRedraw": "GQCH 高清重绘",
    "GQCHDownloadVector": "GQCH 下载矢量图",
    "GQCHEditImage": "GQCH AI改图",
    "GQCHExpandImage": "GQCH 扩图",
    "GQCHRemoveFabric": "GQCH 去布纹",
    "GQCHRemoveDiamonds": "GQCH 去烫钻",
    "GQCHRemoveWatermark": "GQCH 去水印",
}