"""
批量 Nano Banana Pro 节点
ComfyUI 自定义节点，用于批量处理图像生成任务
支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存
"""

import time
import math
import random
import asyncio
import aiohttp
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Tuple, List
from PIL import Image

import torch
import numpy as np

from ..utils.image_utils import tensor_to_pil, pil_to_tensor
from ..utils.file_utils import (
    ImageInfo,
    load_images_from_folder,
    pair_images_indexed,
    pair_images_cartesian,
    generate_output_filename,
    save_image
)
from ..clients.gemini_client import GeminiAPIClient
from ..models_config import get_enabled_models

# 导入 ComfyUI 原生进度条
try:
    from comfy.utils import ProgressBar
    PROGRESS_BAR_AVAILABLE = True
except ImportError:
    PROGRESS_BAR_AVAILABLE = False
    print("⚠️ BatchNanoBananaPro: comfy.utils.ProgressBar 不可用，将只使用终端进度显示")


class BatchNanoBananaPro:
    """
    批量 Nano Banana Pro 节点
    
    功能：
    - 从多个文件夹加载图片
    - 支持三种配对模式：
      * 1:1 - 索引配对（文件夹之间按位置配对）
      * 1*N - 笛卡尔积配对（所有可能组合）
      * 不配对 - 固定参考图模式（文件夹图片依次与所有参考图组合）
    - 批量调用 API 生成图像
    - 智能命名保存（保留原始文件名）
    - 并发控制（默认最大 100）
    
    注意：
    - 「不配对」模式只支持单个文件夹
    - 支持的模型列表从 models_config.py 动态加载
    - 要添加/禁用模型，请编辑 models_config.py 文件
    """
    
    # 支持的模型列表（从配置文件动态加载）
    MODELS = None  # 将在 INPUT_TYPES 中动态获取
    
    # 支持的宽高比列表
    ASPECT_RATIOS = [
        "1:1", "4:3", "3:4", "16:9", "9:16",
        "2:3", "3:2", "4:5", "5:4", "21:9"
    ]
    
    # 支持的分辨率列表
    RESOLUTIONS = ["1K", "2K", "4K"]
    
    # 配对模式
    PAIRING_MODES = ["1:1", "1*N", "不配对"]
    
    def __init__(self):
        """初始化节点"""
        self.client = None
    
    def resize_to_megapixels(
        self,
        image: Image.Image,
        target_megapixels: float
    ) -> Image.Image:
        """
        将图像缩放到指定的总像素数，保持纵横比
        
        Args:
            image: PIL Image 对象
            target_megapixels: 目标像素数（百万像素）
        
        Returns:
            缩放后的 PIL Image
        
        Example:
            >>> resized = self.resize_to_megapixels(img, 2.0)  # 缩放到2百万像素
        """
        # 计算当前像素数
        current_pixels = image.width * image.height
        target_pixels = int(target_megapixels * 1_000_000)
        
        # 如果当前像素数已经接近目标，则不缩放
        if abs(current_pixels - target_pixels) / target_pixels < 0.05:
            return image
        
        # 计算缩放比例
        scale = (target_pixels / current_pixels) ** 0.5
        
        # 计算新尺寸
        new_width = int(image.width * scale)
        new_height = int(image.height * scale)
        
        # 确保至少为1像素
        new_width = max(1, new_width)
        new_height = max(1, new_height)
        
        # 使用 Lanczos 重采样
        resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        
        return resized_image
    
    @classmethod
    def INPUT_TYPES(cls):
        """
        定义输入参数
        
        ComfyUI 节点规范：
        - required: 必选参数
        - optional: 可选参数
        """
        # 从配置文件动态获取启用的模型列表
        enabled_models = get_enabled_models()
        
        # 如果没有启用的模型，使用空列表（会导致节点不可用，提示用户配置）
        if not enabled_models:
            enabled_models = ["请在 models_config.py 中启用至少一个模型"]
        
        # 创建9个独立的图像输入
        optional_inputs = {}
        for i in range(1, 10):  # 1-9
            optional_inputs[f"参考图{i}"] = ("IMAGE",)
        
        return {
            "required": {
                "prompt": ("STRING", {
                    "default": "一个中国女子的OOTD",
                    "multiline": True
                }),
                "模型": (enabled_models, {
                    "default": enabled_models[0]
                }),
                "宽高比": (cls.ASPECT_RATIOS, {
                    "default": "1:1"
                }),
                "分辨率": (cls.RESOLUTIONS, {
                    "default": "2K"
                }),
                "像素缩放": ("BOOLEAN", {
                    "default": False
                }),
                "分辨率像素": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.1,
                    "max": 100.0,
                    "step": 0.1,
                    "display": "number"
                }),
                "seed": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 0xffffffffffffffff
                }),
                "文件夹1": ("STRING", {
                    "default": "",
                    "multiline": False
                }),
                "文件夹2": ("STRING", {
                    "default": "",
                    "multiline": False
                }),
                "文件夹3": ("STRING", {
                    "default": "",
                    "multiline": False
                }),
                "文件夹4": ("STRING", {
                    "default": "",
                    "multiline": False
                }),
                "保存路径": ("STRING", {
                    "default": "",
                    "multiline": False
                }),
                "图片配对模式": (cls.PAIRING_MODES, {
                    "default": "不配对"
                })
            },
            "optional": optional_inputs
        }
    
    # 返回值类型
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("输出图像",)
    
    # 执行函数名
    FUNCTION = "process_batch"
    
    # 节点分类
    CATEGORY = "image/batch"
    
    def _load_folders(
        self,
        folder1: str,
        folder2: Optional[str],
        folder3: Optional[str],
        folder4: Optional[str],
        enable_scaling: bool,
        target_megapixels: float
    ) -> List[List[ImageInfo]]:
        """
        加载所有文件夹中的图片
        
        Args:
            folder1-4: 文件夹路径
            enable_scaling: 是否启用像素缩放
            target_megapixels: 目标像素数（百万像素）
        
        Returns:
            图片列表的列表
        """
        folders = [folder1, folder2, folder3, folder4]
        all_images = []
        
        for i, folder in enumerate(folders, 1):
            if folder and folder.strip():
                try:
                    images = load_images_from_folder(folder)
                    if images:
                        # 应用像素缩放
                        if enable_scaling:
                            scaled_images = []
                            for img_info in images:
                                scaled_img = self.resize_to_megapixels(
                                    img_info.image,
                                    target_megapixels
                                )
                                # 创建新的 ImageInfo，保留其他元数据
                                scaled_info = ImageInfo(
                                    image=scaled_img,
                                    filename=img_info.filename,
                                    extension=img_info.extension,
                                    source_path=img_info.source_path
                                )
                                scaled_images.append(scaled_info)
                            images = scaled_images
                        
                        all_images.append(images)
                        print(f"BatchNanoBananaPro: 文件夹{i} 加载了 {len(images)} 张图片")
                    else:
                        print(f"BatchNanoBananaPro: 文件夹{i} 为空或没有有效图片")
                except ValueError as e:
                    print(f"BatchNanoBananaPro: 文件夹{i} 加载失败 - {e}")
        
        return all_images
    
    def _create_pairs(
        self,
        image_lists: List[List[ImageInfo]],
        pairing_mode: str,
        manual_images: Optional[List[ImageInfo]] = None
    ) -> List[Tuple[ImageInfo, ...]]:
        """
        根据配对模式创建图片组合
        
        Args:
            image_lists: 从文件夹加载的图片列表
            pairing_mode: 配对模式 (1:1, 1*N, 不配对)
            manual_images: 手动输入的参考图
        
        Returns:
            配对后的元组列表
        
        Raises:
            ValueError: 不配对模式下填入多个文件夹时
        """
        # === 新模式：不配对 ===
        if pairing_mode == "不配对":
            # 验证：只支持单个文件夹
            if len(image_lists) > 1:
                raise ValueError("「不配对」模式只支持单个文件夹，请清空其他文件夹路径")
            
            # 场景1：有文件夹 + 有参考图
            if image_lists and manual_images:
                folder_images = image_lists[0]
                # 每张文件夹图片 + 所有参考图
                pairs = []
                for img in folder_images:
                    pair = (img,) + tuple(manual_images)
                    pairs.append(pair)
                return pairs
            
            # 场景2：有文件夹 + 无参考图
            elif image_lists:
                # 每张图片单独成组
                return [(img,) for img in image_lists[0]]
            
            # 场景3：无文件夹 + 有参考图
            elif manual_images:
                # 每张参考图单独成组
                return [(img,) for img in manual_images]
            
            else:
                return []
        
        # === 原有逻辑：1:1 和 1*N ===
        # 如果有手动参考图，添加到列表中（所有参考图作为一个列表）
        if manual_images:
            image_lists.append(manual_images)
        
        if not image_lists:
            return []
        
        # 如果只有一个列表，直接返回每个图片作为单元素元组
        if len(image_lists) == 1:
            return [(img,) for img in image_lists[0]]
        
        # 根据配对模式选择配对函数
        if pairing_mode == "1:1":
            pairs = pair_images_indexed(*image_lists)
        else:  # 1*N
            pairs = pair_images_cartesian(*image_lists)
        
        return pairs
    
    async def _generate_single_task(
        self,
        client: GeminiAPIClient,
        session: aiohttp.ClientSession,
        prompt: str,
        model: str,
        resolution: str,
        aspect_ratio: str,
        images: List[ImageInfo],
        output_folder: str,
        task_index: int
    ) -> dict:
        """
        执行单个生成任务
        
        Args:
            client: API 客户端
            session: aiohttp 会话
            prompt: 提示词
            model: 模型名称
            resolution: 分辨率
            aspect_ratio: 宽高比
            images: 输入图片列表
            output_folder: 输出文件夹
            task_index: 任务索引
        
        Returns:
            包含结果信息的字典
        """
        result = {
            "task_index": task_index,
            "success": False,
            "generated_count": 0,
            "saved_files": [],
            "error": None
        }
        
        try:
            # 准备输入图片
            input_pil_images = [info.image for info in images]
            
            # 调用 API 生成图片（固定生成1次）
            generated_images = []
            try:
                gen_result = await client.generate_single_async(
                    prompt=prompt,
                    model=model,
                    resolution=resolution,
                    aspect_ratio=aspect_ratio,
                    images=input_pil_images,
                    session=session
                )
                if gen_result:
                    generated_images.extend(gen_result)
            except Exception as e:
                error_msg = str(e)
                print(f"BatchNanoBananaPro: 任务 {task_index + 1} 生成失败 - {error_msg}")
                result["error"] = error_msg
            
            # 保存生成的图片
            for i, gen_img in enumerate(generated_images):
                # 使用任务索引作为唯一标识，确保并发安全
                output_path = generate_output_filename(
                    source_images=list(images),
                    batch_index=i,
                    output_folder=output_folder,
                    extension=".png",
                    task_id=f"task{task_index}"
                )
                save_image(gen_img, output_path)
                result["saved_files"].append(output_path)
            
            # 只有生成了图片才标记为成功
            if len(generated_images) > 0:
                result["success"] = True
            result["generated_count"] = len(generated_images)
            
        except Exception as e:
            result["error"] = str(e)
        
        return result
    
    async def _process_batch_async(
        self,
        pairs: List[Tuple[ImageInfo, ...]],
        prompt: str,
        model: str,
        resolution: str,
        aspect_ratio: str,
        output_folder: str,
        pbar=None
    ) -> List[dict]:
        """
        异步批量处理所有任务
        
        Args:
            pairs: 配对后的图片组合
            prompt: 提示词
            model: 模型名称
            resolution: 分辨率
            aspect_ratio: 宽高比
            output_folder: 输出文件夹
        
        Returns:
            所有任务的结果列表
        """
        if self.client is None:
            self.client = GeminiAPIClient()
        
        # 固定最大并发数为 100
        max_concurrent = 100
        
        total_tasks = len(pairs)
        all_results = []
        completed = 0
        success_count = 0
        fail_count = 0
        
        # 计算分批数量
        num_batches = math.ceil(total_tasks / max_concurrent)
        
        # 进度打印配置：任务数 >= 50 时，额外显示百分比里程碑
        show_milestone = total_tasks >= 50
        milestones = [0.2, 0.4, 0.6, 0.8, 1.0]  # 20%, 40%, 60%, 80%, 100%
        milestone_index = 0
        
        if num_batches > 1:
            print(f"BatchNanoBananaPro: 任务数 {total_tasks} 超过并发上限 {max_concurrent}，将分 {num_batches} 批执行")
        
        connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
        
        async with aiohttp.ClientSession(connector=connector) as session:
            for batch_idx in range(num_batches):
                start_idx = batch_idx * max_concurrent
                end_idx = min(start_idx + max_concurrent, total_tasks)
                batch_pairs = pairs[start_idx:end_idx]
                
                if num_batches > 1:
                    print(f"BatchNanoBananaPro: 执行第 {batch_idx + 1}/{num_batches} 批 ({start_idx + 1}-{end_idx})...")
                
                # 创建当前批次的任务
                tasks = []
                for i, pair in enumerate(batch_pairs):
                    task = asyncio.create_task(
                        self._generate_single_task(
                            client=self.client,
                            session=session,
                            prompt=prompt,
                            model=model,
                            resolution=resolution,
                            aspect_ratio=aspect_ratio,
                            images=list(pair),
                            output_folder=output_folder,
                            task_index=start_idx + i
                        )
                    )
                    tasks.append(task)
                
                # 使用 as_completed 实时获取完成的任务
                for coro in asyncio.as_completed(tasks):
                    result_data = None
                    try:
                        result = await coro
                        if isinstance(result, Exception):
                            result_data = {
                                "success": False,
                                "error": str(result),
                                "generated_count": 0,
                                "saved_files": []
                            }
                            all_results.append(result_data)
                        else:
                            result_data = result
                            all_results.append(result)
                    except Exception as e:
                        result_data = {
                            "success": False,
                            "error": str(e),
                            "generated_count": 0,
                            "saved_files": []
                        }
                        all_results.append(result_data)
                    
                    completed += 1
                    
                    # 根据成功/失败状态打印不同信息
                    if result_data and result_data.get("success", False):
                        success_count += 1
                        print(f"BatchNanoBananaPro: 任务 {completed}/{total_tasks} 成功 ✓")
                    else:
                        fail_count += 1
                        # 提取错误信息的第一行
                        error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
                        # 截取第一行或前50个字符
                        if '\n' in error_msg:
                            error_msg = error_msg.split('\n')[0]
                        if len(error_msg) > 50:
                            error_msg = error_msg[:50] + "..."
                        print(f"BatchNanoBananaPro: 任务 {completed}/{total_tasks} 失败 ✗ - {error_msg}")
                    
                    # 更新 ComfyUI 原生进度条
                    if pbar is not None:
                        pbar.update(1)
                    
                    # 大任务额外显示百分比里程碑
                    if show_milestone and milestone_index < len(milestones):
                        progress = completed / total_tasks
                        if progress >= milestones[milestone_index]:
                            percentage = int(milestones[milestone_index] * 100)
                            print(f"BatchNanoBananaPro: >>> 进度 {percentage}% <<<")
                            milestone_index += 1
        
        return all_results
    
    def process_batch(
        self,
        prompt: str,
        文件夹1: str,
        文件夹2: str,
        文件夹3: str,
        文件夹4: str,
        像素缩放: bool,
        分辨率像素: float,
        seed: int,
        保存路径: str,
        图片配对模式: str,
        模型: str,
        宽高比: str,
        分辨率: str,
        **kwargs
    ) -> Tuple[torch.Tensor]:
        """
        批量处理图像生成任务
        
        Args:
            prompt: 提示词
            文件夹1-4: 图片文件夹路径
            像素缩放: 是否启用像素缩放
            分辨率像素: 目标像素数（百万像素）
            seed: 随机种子
            保存路径: 输出保存路径
            图片配对模式: 1:1 或 1*N
            模型: 模型名称
            宽高比: 输出宽高比
            分辨率: 输出分辨率
            **kwargs: 动态参考图输入 (参考图1-9)
        
        Returns:
            输出图像张量
        """
        start_time = time.time()
        
        try:
            # 设置随机种子（用于本地随机操作）
            random.seed(seed)
            np.random.seed(seed % (2**32))
            # 验证保存路径
            if not 保存路径 or not 保存路径.strip():
                raise ValueError("请提供保存路径")
            
            # 加载文件夹图片
            print("BatchNanoBananaPro: 开始加载图片...")
            image_lists = self._load_folders(
                文件夹1, 文件夹2, 文件夹3, 文件夹4,
                像素缩放, 分辨率像素
            )
            
            # 处理独立的参考图输入
            manual_images = []
            for i in range(1, 10):  # 1-9
                key = f"参考图{i}"
                if key in kwargs and kwargs[key] is not None:
                    pil_images = tensor_to_pil(kwargs[key])
                    for j, img in enumerate(pil_images):
                        # 如果启用像素缩放，也对参考图进行缩放
                        if 像素缩放:
                            img = self.resize_to_megapixels(img, 分辨率像素)
                        
                        manual_images.append(
                            ImageInfo(
                                image=img,
                                filename=f"manual_{i}_{j}",
                                extension=".png",
                                source_path=""
                            )
                        )
            
            if manual_images:
                print(f"BatchNanoBananaPro: 加载了 {len(manual_images)} 张参考图")
            
            # 验证是否有图片
            total_folder_images = sum(len(lst) for lst in image_lists)
            total_manual_images = len(manual_images)
            
            if total_folder_images == 0 and total_manual_images == 0:
                raise ValueError("未找到任何图片，请检查文件夹路径或提供参考图")
            
            # 创建配对
            print(f"BatchNanoBananaPro: 使用 {图片配对模式} 模式创建配对...")
            pairs = self._create_pairs(image_lists, 图片配对模式, manual_images if manual_images else None)
            
            if not pairs:
                raise ValueError("配对结果为空，请检查输入")
            
            total_tasks = len(pairs)
            print(f"BatchNanoBananaPro: 共 {total_tasks} 组配对")
            
            # 创建 ComfyUI 原生进度条
            pbar = None
            if PROGRESS_BAR_AVAILABLE:
                pbar = ProgressBar(total_tasks)
            
            # 初始化 API 客户端
            if self.client is None:
                try:
                    self.client = GeminiAPIClient()
                except ValueError as e:
                    raise ValueError(f"初始化 API 客户端失败: {str(e)}")
            
            # 执行批量生成
            print("BatchNanoBananaPro: 开始批量生成...")
            
            # 在新线程中运行异步代码，避免事件循环冲突
            def run_async_in_thread():
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                try:
                    return loop.run_until_complete(
                        self._process_batch_async(
                            pairs=pairs,
                            prompt=prompt,
                            model=模型,
                            resolution=分辨率,
                            aspect_ratio=宽高比,
                            output_folder=保存路径,
                            pbar=pbar
                        )
                    )
                finally:
                    loop.close()
            
            # 使用线程池在新线程中运行事件循环
            with ThreadPoolExecutor(max_workers=1) as executor:
                future = executor.submit(run_async_in_thread)
                results = future.result()
            
            # 统计结果
            success_count = sum(1 for r in results if r.get("success", False))
            fail_count = len(results) - success_count
            total_generated = sum(r.get("generated_count", 0) for r in results)
            all_saved_files = []
            for r in results:
                all_saved_files.extend(r.get("saved_files", []))
            
            elapsed = time.time() - start_time
            
            # 精简统计信息
            print("=" * 50)
            print(f"BatchNanoBananaPro 处理完成 | 总耗时: {elapsed:.2f}s | 成功: {success_count}/{total_tasks} | 生成: {total_generated}张")
            print(f"保存路径: {保存路径}")
            
            # 失败详情（如果有）
            failed_results = [r for r in results if not r.get("success", False)]
            if failed_results:
                # 收集失败任务的索引
                failed_indices = [str(r.get('task_index', '?') + 1) for r in failed_results[:5]]
                failed_str = ",".join(failed_indices)
                if len(failed_results) > 5:
                    failed_str += f"... (共{len(failed_results)}个)"
                # 显示第一个失败原因作为示例
                first_error = failed_results[0].get('error', '未知错误')
                print(f"失败 {len(failed_results)}个: 任务{failed_str} - {first_error}")
            
            # 收集所有生成的图片
            output_images = []
            for file_path in all_saved_files:
                try:
                    img = Image.open(file_path)
                    output_images.append(img)
                except Exception as e:
                    print(f"BatchNanoBananaPro: 无法加载图片 {file_path} - {e}")
            
            # 如果没有生成成功的图片，创建一个占位图
            if not output_images:
                placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
                output_images = [placeholder]
            
            # 转换为张量
            output_tensor = pil_to_tensor(output_images)
            
            return (output_tensor,)
        
        except ValueError as e:
            # 检测是否为授权错误
            if str(e) == "未授权！":
                print("请联系作者授权后方可使用！")
            else:
                print(f"BatchNanoBananaPro: 输入错误 - {str(e)}")
            raise
        
        except RuntimeError as e:
            print(f"BatchNanoBananaPro: 运行时错误 - {str(e)}")
            raise
        
        except Exception as e:
            print(f"BatchNanoBananaPro: 未知错误 - {str(e)}")
            raise
        
        finally:
            # 无论成功或失败，都尝试查询余额
            if self.client is not None:
                try:
                    balance_data = self.client.query_balance_sync()
                    balance_info = self.client.format_balance_info(balance_data)
                    print(f"{balance_info}")
                    print("=" * 50)
                except Exception as e:
                    print(f"⚠️ 余额查询失败 - {str(e)}")
                    print("=" * 50)