"""
Nano Banana 节点 (V3)
ComfyUI 自定义节点，用于调用生图模型（OpenAI 兼容接口）
使用 V3 DynamicCombo 实现模型-宽高比-分辨率动态联动
"""

import io as _io
import re
import json
import time
import math
import base64
import random
import asyncio
import aiohttp
from concurrent.futures import ThreadPoolExecutor
from typing import List, Optional

import torch
import numpy as np
from PIL import Image

from comfy_api.latest import io

from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts, encode_image_to_base64, encode_image_to_base64_limited
from ..utils.config import (
    NETWORK_ROUTE_OPTIONS,
    get_base_url_by_route,
    get_api_key_or_raise,
)
from ..utils.http_error import RETRYABLE_STATUS_CODES, HTTP_ERROR_MESSAGES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR
from ..clients.gemini_client import GeminiAPIClient

try:
    import folder_paths
    FOLDER_PATHS_AVAILABLE = True
except ImportError:
    FOLDER_PATHS_AVAILABLE = False

try:
    from comfy.utils import ProgressBar
    PROGRESS_BAR_AVAILABLE = True
except ImportError:
    PROGRESS_BAR_AVAILABLE = False

try:
    import psutil
    MEMORY_MONITOR_AVAILABLE = True
except ImportError:
    MEMORY_MONITOR_AVAILABLE = False

DEBUG_LOG_ENABLED = True
REQUEST_LOG_ENABLED = False

_NODE = "Nano Banana"
_ENDPOINT = "/v1/chat/completions"

_client_instance = None


def _get_client():
    global _client_instance
    if _client_instance is None:
        _client_instance = GeminiAPIClient()
    return _client_instance


def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
    if not images:
        placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
        return pil_to_tensor([placeholder])

    base_size = max(images, key=lambda img: img.size[0] * img.size[1]).size
    matched = [img for img in images if img.size == base_size]
    skipped = [img for img in images if img.size != base_size]

    if skipped:
        sizes_str = ", ".join(f"{img.size[0]}x{img.size[1]}" for img in skipped)
        print(
            f"{node_label}: 丢弃 {len(skipped)} 张较小尺寸的图 ({sizes_str})，"
            f"仅输出最大尺寸 {base_size[0]}x{base_size[1]} 的 {len(matched)} 张"
        )

    return pil_to_tensor(matched)


MODEL_ID_MAP = {
    "Nano Banana Pro": "nano-banana-pro",
    "Nano Banana 2": "nano-banana-2",
    "Nano Banana": "nano-banana",
}
RESOLUTION_KEY_MAP = {
    "512px": "0.5k",
    "1K": "1k",
    "2K": "2k",
    "4K": "4k",
}
BILLING_SPECIAL_ONLY = {"nano-banana"}


def _build_model_id(model_name: str, resolution: str, billing: str) -> str:
    base = MODEL_ID_MAP.get(model_name, "nano-banana-pro")

    if base == "nano-banana":
        if billing == "官方":
            raise ValueError(f"模型 \"{model_name}\" 仅支持特价计费")
        return "nano-banana"

    res_key = RESOLUTION_KEY_MAP.get(resolution, "2k")
    is_official = (billing == "官方")

    if base == "nano-banana-pro" and res_key == "1k" and not is_official:
        return "nano-banana-pro"

    if base == "nano-banana-2" and res_key == "0.5k":
        if is_official:
            raise ValueError("Nano Banana 2 的 512px 分辨率仅支持特价计费")
        return "nano-banana-2-0.5k"

    model_id = f"{base}-{res_key}"
    if is_official:
        model_id += "-official"
    return model_id

_IMAGE_RE = re.compile(r"!\[.*?\]\(data:image/(\w+);base64,([A-Za-z0-9+/=]+)\)")


def _get_headers(api_key: str) -> dict:
    return {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "X-Accel-Buffering": "no",
        "Cache-Control": "no-cache, no-transform",
    }


def _build_request_body(
    prompt: str,
    model: str,
    aspect_ratio: str,
    resolution: str,
    images: Optional[List[Image.Image]] = None,
    enable_grounding: bool = False,
    thinking_level: Optional[str] = None,
) -> dict:
    content_parts = [{"type": "text", "text": prompt}]

    if images:
        for img in images:
            b64 = encode_image_to_base64_limited(img, format="PNG")
            content_parts.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/png;base64,{b64}"}
            })

    body = {
        "model": model,
        "stream": True,
        "messages": [{"role": "user", "content": content_parts}],
    }

    google_config = {
        "image_config": {
            "image_size": resolution,
        }
    }
    if aspect_ratio and aspect_ratio != "智能":
        google_config["image_config"]["aspect_ratio"] = aspect_ratio
    if thinking_level:
        google_config["thinking_config"] = {
            "thinking_level": thinking_level.lower(),
            "include_thoughts": True,
        }
    body["extra_body"] = {"google": google_config}

    if enable_grounding:
        body["extra_body"]["google_search"] = True

    return body


async def _generate_single(
    session: aiohttp.ClientSession,
    base_url: str,
    api_key: str,
    prompt: str,
    model: str,
    resolution: str,
    aspect_ratio: str,
    images: Optional[List[Image.Image]] = None,
    enable_grounding: bool = False,
    thinking_level: Optional[str] = None,
) -> List[Image.Image]:
    url = f"{base_url}{_ENDPOINT}"
    headers = _get_headers(api_key)
    body = _build_request_body(
        prompt=prompt,
        model=model,
        aspect_ratio=aspect_ratio,
        resolution=resolution,
        images=images,
        enable_grounding=enable_grounding,
        thinking_level=thinking_level,
    )

    if REQUEST_LOG_ENABLED:
        extra = json.dumps(body.get("extra_body", {}), ensure_ascii=False)
        print(f"[请求] POST {url} | model={model} | extra_body={extra}")

    last_status = None
    for attempt in range(DEFAULT_MAX_RETRIES + 1):
        resp = await session.post(url, headers=headers, json=body)
        if resp.status == 200:
            break
        last_status = resp.status
        if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
            friendly = HTTP_ERROR_MESSAGES.get(resp.status, f"请求失败 ({resp.status})")
            delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
            print(f"Nano Banana: {friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
            resp.close()
            await asyncio.sleep(delay)
            continue
        error_text = await resp.text()
        resp.close()
        if resp.status in HTTP_ERROR_MESSAGES:
            raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
        try:
            err_json = json.loads(error_text)
            msg = err_json.get("error", {}).get("message", error_text[:200])
        except Exception:
            msg = error_text[:200]
        raise RuntimeError(f"API 错误 ({resp.status}): {msg}")
    else:
        if last_status and last_status in HTTP_ERROR_MESSAGES:
            raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
        raise RuntimeError(f"API 错误: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")

    full_content = ""
    buffer = ""
    t_request = time.time()
    t_first_token = None
    async for raw_chunk in resp.content.iter_any():
        if t_first_token is None:
            t_first_token = time.time()
        buffer += raw_chunk.decode("utf-8")
        while "\n" in buffer:
            line_str, buffer = buffer.split("\n", 1)
            line_str = line_str.strip()
            if not line_str or not line_str.startswith("data:"):
                continue
            data_str = line_str[5:].strip()
            if data_str == "[DONE]":
                break
            try:
                chunk = json.loads(data_str)
                delta = chunk.get("choices", [{}])[0].get("delta", {})
                if "content" in delta:
                    full_content += delta["content"]
            except (json.JSONDecodeError, IndexError):
                continue
    t_done = time.time()
    resp.close()

    if not full_content:
        raise RuntimeError("API 未返回有效内容")

    # 思考模型可能输出多张临时图片，最终图片始终是最后一张
    matches = list(_IMAGE_RE.finditer(full_content))
    if not matches:
        raise RuntimeError(f"响应中未找到图片: {full_content[:100]}")

    last_match = matches[-1]
    img_data = base64.b64decode(last_match.group(2))
    final_image = Image.open(_io.BytesIO(img_data)).convert("RGB")

    first_token_ms = (t_first_token - t_request) * 1000 if t_first_token else 0
    download_ms = (t_done - t_first_token) * 1000 if t_first_token else 0

    return [final_image], first_token_ms, download_ms


async def _generate_single_task(
    session: aiohttp.ClientSession,
    base_url: str,
    api_key: str,
    prompt: str,
    model: str,
    resolution: str,
    aspect_ratio: str,
    images: Optional[List[Image.Image]],
    global_task_index: int,
    enable_grounding: bool = False,
    thinking_level: Optional[str] = None,
) -> dict:
    result = {
        "global_task_index": global_task_index,
        "prompt": prompt,
        "success": False,
        "generated_count": 0,
        "output_images": [],
        "error": None,
    }
    try:
        gen_images, first_token_ms, download_ms = await _generate_single(
            session=session,
            base_url=base_url,
            api_key=api_key,
            prompt=prompt,
            model=model,
            resolution=resolution,
            aspect_ratio=aspect_ratio,
            images=images if images else None,
            enable_grounding=enable_grounding,
            thinking_level=thinking_level,
        )
        result["output_images"] = gen_images
        result["success"] = True
        result["generated_count"] = len(gen_images)
    except Exception as e:
        result["error"] = str(e)
    return result


async def _process_batch_async(
    base_url: str,
    api_key: str,
    prompts: List[str],
    model: str,
    resolution: str,
    aspect_ratio: str,
    images_per_prompt: int,
    input_images: Optional[List[Image.Image]],
    pbar=None,
    enable_grounding: bool = False,
    thinking_level: Optional[str] = None,
) -> List[dict]:
    tasks_def = []
    for p_idx, prompt in enumerate(prompts):
        for sub_idx in range(images_per_prompt):
            tasks_def.append((p_idx, sub_idx, prompt))

    total_tasks = len(tasks_def)
    max_concurrent = 50
    num_batches = math.ceil(total_tasks / max_concurrent)

    all_results = []
    completed = 0
    success_count = 0
    fail_count = 0

    connector = aiohttp.TCPConnector(ssl=False, 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)

            tasks = []
            for i in range(start_idx, end_idx):
                _, _, prompt = tasks_def[i]
                task = asyncio.create_task(
                    _generate_single_task(
                        session=session,
                        base_url=base_url,
                        api_key=api_key,
                        prompt=prompt,
                        model=model,
                        resolution=resolution,
                        aspect_ratio=aspect_ratio,
                        images=input_images,
                        global_task_index=i,
                        enable_grounding=enable_grounding,
                        thinking_level=thinking_level,
                    )
                )
                tasks.append(task)

            batch_results = []
            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, "output_images": [], "prompt": ""}
                    else:
                        result_data = result
                except Exception as e:
                    result_data = {"success": False, "error": str(e), "generated_count": 0, "output_images": [], "prompt": ""}

                batch_results.append(result_data)
                completed += 1
                prompt_snippet = (result_data.get("prompt", "") or "")[:30]

                if result_data and result_data.get("success", False):
                    success_count += 1
                    count = result_data.get("generated_count", 1)
                    print(f"Nano Banana: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✓成功({count}张)")
                else:
                    fail_count += 1
                    error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
                    print(f"Nano Banana: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}")

                if pbar is not None:
                    pbar.update(1)

            all_results.extend(batch_results)
            import gc; gc.collect()
            await asyncio.sleep(0.1)

    return all_results


class NanoBanana(io.ComfyNode):

    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="NanoBanana",
            display_name="Nano Banana",
            category="image/generation",
            inputs=[
                io.String.Input(
                    "prompt",
                    default="一个中国女子的OOTD",
                    multiline=True,
                ),
                io.DynamicCombo.Input("模型", options=[
                    io.DynamicCombo.Option("Nano Banana Pro", [
                        io.Combo.Input("宽高比", options=[
                            "智能", "1:1", "2:3", "3:2", "3:4", "4:3",
                            "4:5", "5:4", "9:16", "16:9", "21:9",
                        ], default="智能"),
                        io.Combo.Input("分辨率", options=["1K", "2K", "4K"], default="2K"),
                    ]),
                    io.DynamicCombo.Option("Nano Banana 2", [
                        io.Combo.Input("宽高比", options=[
                            "智能", "1:1", "1:4", "1:8", "2:3", "3:2", "3:4",
                            "4:1", "4:3", "4:5", "5:4", "8:1",
                            "9:16", "16:9", "21:9",
                        ], default="智能"),
                        io.Combo.Input("分辨率", options=["512px", "1K", "2K", "4K"], default="2K"),
                        io.Combo.Input("思考深度", options=["高", "低"], default="高"),
                    ]),
                    io.DynamicCombo.Option("Nano Banana", [
                        io.Combo.Input("宽高比", options=[
                            "智能", "1:1", "2:3", "3:2", "3:4", "4:3",
                            "4:5", "5:4", "9:16", "16:9", "21:9",
                        ], default="智能"),
                        io.Combo.Input("分辨率", options=["1K"], default="1K"),
                    ]),
                ]),
                io.Int.Input("生图数量", default=1, min=1, max=1000, step=1),
                io.Combo.Input("网络", options=NETWORK_ROUTE_OPTIONS, default="全球加速"),
                io.Combo.Input("计费", options=["特价", "官方"], default="特价"),
                io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"),
                io.Int.Input("seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF),
                io.Image.Input("参考图1", optional=True),
                io.Image.Input("参考图2", optional=True),
                io.Image.Input("参考图3", optional=True),
                io.Image.Input("参考图4", optional=True),
                io.Image.Input("参考图5", optional=True),
                io.Image.Input("参考图6", optional=True),
                io.Image.Input("参考图7", optional=True),
                io.Image.Input("参考图8", optional=True),
                io.Image.Input("参考图9", optional=True),
            ],
            outputs=[
                io.Image.Output(display_name="输出图像"),
            ],
        )

    @classmethod
    def execute(cls, prompt, 模型, 生图数量, 计费, 网络, 谷歌搜索, seed, **kwargs) -> io.NodeOutput:
        start_time = time.time()

        model_name = 模型["模型"]
        宽高比 = 模型["宽高比"]
        分辨率 = 模型["分辨率"]
        思考深度 = 模型.get("思考深度")

        enable_grounding = (谷歌搜索 == "打开")

        thinking_level = None
        if model_name == "Nano Banana 2" and 思考深度:
            thinking_level = "High" if 思考深度 == "高" else "Low"

        actual_model = _build_model_id(model_name, 分辨率, 计费)

        api_key = get_api_key_or_raise("O1KEY_API_KEY")
        base_url = get_base_url_by_route(网络)

        pbar = ProgressBar(生图数量) if PROGRESS_BAR_AVAILABLE else None

        try:
            random.seed(seed)
            np.random.seed(seed % (2**32))

            input_images = []
            for i in range(1, 10):
                key = f"参考图{i}"
                if key in kwargs and kwargs[key] is not None:
                    pil_imgs = tensor_to_pil(kwargs[key])
                    input_images.extend(pil_imgs)

            if input_images and len(input_images) > 14:
                raise ValueError(f"输入图像数量 {len(input_images)} 超过限制 14 张")

            batch_prompts = parse_batch_prompts(prompt)

            grounding_str = " | 谷歌搜索接地" if enable_grounding else ""
            thinking_str = f" | 思考:{thinking_level}" if thinking_level else ""
            if batch_prompts:
                num_prompts = len(batch_prompts)
                total_images = num_prompts * 生图数量
                mode_str = f"批量提示词模式 ({num_prompts}个提示词)"
                if input_images:
                    mode_str += f" (输入{len(input_images)}张)"
                print(f"Nano Banana: {mode_str} | {分辨率} {宽高比} | 共{total_images}张{grounding_str}{thinking_str}")
            else:
                mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式"
                print(f"Nano Banana: {mode_str} | {分辨率} {宽高比} | {生图数量}张{grounding_str}{thinking_str}")

            if batch_prompts or 生图数量 > 1:
                prompts = batch_prompts if batch_prompts else [prompt]
                images_per_prompt = 生图数量
                total_tasks = len(prompts) * images_per_prompt

                if pbar is not None:
                    pbar = ProgressBar(total_tasks)

                def run_async_in_thread():
                    loop = asyncio.new_event_loop()
                    asyncio.set_event_loop(loop)
                    try:
                        return loop.run_until_complete(
                            _process_batch_async(
                                base_url=base_url,
                                api_key=api_key,
                                prompts=prompts,
                                model=actual_model,
                                resolution=分辨率,
                                aspect_ratio=宽高比,
                                images_per_prompt=images_per_prompt,
                                input_images=input_images,
                                pbar=pbar,
                                enable_grounding=enable_grounding,
                                thinking_level=thinking_level,
                            )
                        )
                    finally:
                        loop.close()

                with ThreadPoolExecutor(max_workers=1) as executor:
                    future = executor.submit(run_async_in_thread)
                    try:
                        results = future.result(timeout=900)
                    except TimeoutError:
                        raise RuntimeError("任务执行超时（900秒）")

                success_count = sum(1 for r in results if r.get("success", False))
                fail_count = len(results) - success_count
                elapsed = time.time() - start_time
                time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
                print(f"完成！总耗时 {time_str} | 成功: {success_count}/{total_tasks} | 失败: {fail_count}")

                output_images = []
                for r in results:
                    output_images.extend(r.get("output_images", []))
                if not output_images:
                    placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
                    output_images = [placeholder]

                output_tensor = _images_to_tensor_safe(output_images, _NODE)
                import gc; gc.collect()
                return io.NodeOutput(output_tensor)

            else:
                def run_single():
                    loop = asyncio.new_event_loop()
                    asyncio.set_event_loop(loop)
                    try:
                        async def _do():
                            connector = aiohttp.TCPConnector(ssl=False)
                            async with aiohttp.ClientSession(connector=connector) as session:
                                return await _generate_single(
                                    session=session,
                                    base_url=base_url,
                                    api_key=api_key,
                                    prompt=prompt,
                                    model=actual_model,
                                    resolution=分辨率,
                                    aspect_ratio=宽高比,
                                    images=input_images if input_images else None,
                                    enable_grounding=enable_grounding,
                                    thinking_level=thinking_level,
                                )
                        return loop.run_until_complete(_do())
                    finally:
                        loop.close()

                with ThreadPoolExecutor(max_workers=1) as executor:
                    future = executor.submit(run_single)
                    generated_images, first_token_ms, download_ms = future.result(timeout=900)

                if pbar is not None:
                    pbar.update(1)

                output_tensor = _images_to_tensor_safe(generated_images, _NODE)
                elapsed = time.time() - start_time
                time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
                ft_str = f"{first_token_ms/1000:.2f}s"
                dl_str = f"{download_ms/1000:.2f}s"
                print(f"完成！总耗时 {time_str} | 首字 {ft_str} | 下载 {dl_str} | 成功 {len(generated_images)}张")

                import gc; gc.collect()
                return io.NodeOutput(output_tensor)

        except ValueError as e:
            if str(e) == "未授权！":
                print("请联系作者授权后方可使用！")
                raise ValueError("未授权！") from None
            raise ValueError(str(e)) from None
        except RuntimeError as e:
            raise RuntimeError(str(e)) from None
        except Exception as e:
            raise type(e)(str(e)) from None
        finally:
            try:
                client = _get_client()
                client.base_url = base_url
                balance_data = client.query_balance_sync()
                balance_info = client.format_balance_info(balance_data)
                print(f"Nano Banana: {balance_info}")
            except Exception:
                pass
            import gc; gc.collect()