"""
多路径文件批量加载节点。

用途：一个节点解析用户传入的字符串（多种分隔符），按类型（图片/视频/音频）
分别输出到固定数量的端口。下游若某个端口无输入需要，直接不连线即可 ——
本节点对"没有对应文件"的端口返回 None，ComfyUI 不会执行未连线的下游，
从而保证一个工作流适配 0~N 个素材。

默认端口数量与 MiniMaxH3ReferenceToVideo 对齐：9 图 / 3 视频 / 3 音频。
"""

import os
import re
import subprocess
import tempfile

# 重量级依赖延迟到函数内 import，方便离线做纯字符串逻辑的单测
try:
    import folder_paths  # ComfyUI 运行环境注入
except Exception:  # 允许脱离 ComfyUI 时被静态解析
    folder_paths = None


# 与 MiniMax H3 对齐的默认端口数量
MAX_IMAGES = 9
MAX_VIDEOS = 3
MAX_AUDIOS = 3

# 支持的扩展名（判断类型用；用户显式把文件填到哪个字段就走哪个字段）
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tif", ".tiff"}
VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".avi", ".webm", ".flv", ".m4v"}
AUDIO_EXTS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg", ".opus"}

# 拆分正则：逗号/分号/竖线/花括号/中英文引号/换行/制表符/多空格
# 反斜杠 \ 也作为分隔符（用户描述里提到），但注意 Windows 路径本身含 \，
# 因此优先尝试解析为原始路径，只有当整体不是合法文件时才用 \ 拆一次。
_SPLIT_PATTERN = re.compile(r"[,;|\n\r\t{}\[\]\"\u201c\u201d\u2018\u2019']+")
# 看起来像 Windows 绝对路径（C:\..., D:/...）
_WIN_ABS = re.compile(r"^[A-Za-z]:[\\/]")


def _looks_absolute(p):
    """跨平台的"看起来是绝对路径"判断：兼顾 POSIX 和 Windows。"""
    return os.path.isabs(p) or bool(_WIN_ABS.match(p))


def _split_paths(text):
    """把一段用户输入的字符串切成候选路径列表。

    规则：
    1. 先按正则拆一遍
    2. 每段 strip 掉空白和成对包裹的引号
    3. 单个段落如果不是已存在的文件，再尝试按空格/反斜杠二次拆分
       （这样 "C:\\a.png" 这种 Windows 绝对路径整体不会被误拆）
    """
    if not text:
        return []
    raw_parts = _SPLIT_PATTERN.split(text)
    out = []
    for p in raw_parts:
        p = p.strip().strip("\"'\u201c\u201d\u2018\u2019")
        if not p:
            continue
        if os.path.isfile(p) or _looks_absolute(p):
            # 整段就是合法文件（或至少是绝对路径），直接收
            out.append(p)
            continue
        # 二次拆：空格和反斜杠。到这里说明段里可能塞了多个相对路径
        for q in re.split(r"[\s\\]+", p):
            q = q.strip()
            if q:
                out.append(q)
    return out


def _resolve(path):
    """把路径解析成一个真实存在的文件路径。找不到返回 None。

    顺序：绝对路径 → ComfyUI input 目录相对路径。
    """
    if not path:
        return None
    if os.path.isfile(path):
        return path
    if folder_paths is not None:
        candidate = os.path.join(folder_paths.get_input_directory(), path)
        if os.path.isfile(candidate):
            return candidate
    return None


def _guess_kind(path):
    ext = os.path.splitext(path)[1].lower()
    if ext in IMAGE_EXTS:
        return "image"
    if ext in VIDEO_EXTS:
        return "video"
    if ext in AUDIO_EXTS:
        return "audio"
    return None


# ---------- 加载器 ----------

def _load_image(path):
    """PIL 加载 → (1, H, W, C) float32 in [0,1]。"""
    import numpy as np
    import torch
    from PIL import Image, ImageOps
    with Image.open(path) as im:
        im = ImageOps.exif_transpose(im)
        if im.mode == "RGBA":
            im = im.convert("RGBA")
        elif im.mode != "RGB":
            im = im.convert("RGB")
        arr = np.array(im).astype(np.float32) / 255.0
    if arr.ndim == 2:
        arr = np.expand_dims(arr, -1)
    tensor = torch.from_numpy(arr).unsqueeze(0)
    return tensor


def _extract_audio(video_path):
    """用 ffmpeg 抽视频音轨为 wav 再读入。没有音轨返回静音 dict。"""
    import torch
    sample_rate = 44100
    silent = {
        "waveform": torch.zeros((1, 2, 1), dtype=torch.float32),
        "sample_rate": sample_rate,
    }
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    tmp_path = tmp.name
    tmp.close()
    try:
        cmd = [
            "ffmpeg", "-y", "-i", video_path,
            "-vn", "-acodec", "pcm_s16le", "-ar", str(sample_rate), "-ac", "2",
            "-f", "wav", tmp_path,
        ]
        subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if not os.path.exists(tmp_path) or os.path.getsize(tmp_path) == 0:
            return silent
        import torchaudio
        waveform, sr = torchaudio.load(tmp_path)
        if waveform.dim() == 2:
            waveform = waveform.unsqueeze(0)
        return {"waveform": waveform, "sample_rate": sr}
    except Exception as e:
        print(f"[MultiFileLoader] 抽取音轨失败，返回静音：{e}")
        return silent
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass


def _load_video(path):
    """加载视频为 (T, H, W, C) 帧序列 + 音轨 AUDIO dict。"""
    import cv2
    import numpy as np
    import torch
    cap = cv2.VideoCapture(path)
    if not cap.isOpened():
        raise FileNotFoundError(f"无法打开视频文件: {path}")
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(torch.from_numpy(frame.astype(np.float32) / 255.0))
    cap.release()
    if not frames:
        raise ValueError(f"视频读取到 0 帧: {path}")
    return torch.stack(frames, dim=0), _extract_audio(path)


def _load_audio(path):
    """加载独立音频文件为 ComfyUI AUDIO dict。"""
    import torchaudio
    waveform, sr = torchaudio.load(path)
    if waveform.dim() == 2:
        waveform = waveform.unsqueeze(0)
    return {"waveform": waveform, "sample_rate": sr}


# ---------- 节点主体 ----------

class MultiFileLoader:
    """一次解析多组路径字符串，按类型分发到固定数量的输出端口。

    - 未使用的端口返回 None，下游不连线即安全（ComfyUI 不会执行未连线的下游）。
    - 视频端口成对：video_k 输出帧序列 IMAGE，video_audio_k 输出该视频的音轨 AUDIO。
      若视频无音轨则输出静音 dict，避免下游 AudioLoader 报错。
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "optional": {
                "image_paths": ("STRING", {
                    "multiline": True, "default": "",
                    "tooltip": "图片路径，支持 , ; | 换行 空格 引号 花括号 等分隔",
                }),
                "video_paths": ("STRING", {
                    "multiline": True, "default": "",
                    "tooltip": "视频路径，同上分隔规则",
                }),
                "audio_paths": ("STRING", {
                    "multiline": True, "default": "",
                    "tooltip": "音频路径，同上分隔规则",
                }),
                "auto_classify": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "开启时会按扩展名自动纠正类型（比如把 mp4 从 image_paths 归到视频输出）",
                }),
            }
        }

    # 动态生成 RETURN_TYPES/RETURN_NAMES
    RETURN_TYPES = (
        *(("IMAGE",) * MAX_IMAGES),
        *(("IMAGE", "AUDIO") * MAX_VIDEOS),
        *(("AUDIO",) * MAX_AUDIOS),
        "INT", "INT", "INT",
    )
    RETURN_NAMES = (
        *(f"image_{i + 1}" for i in range(MAX_IMAGES)),
        *(name for i in range(MAX_VIDEOS) for name in (f"video_{i + 1}", f"video_audio_{i + 1}")),
        *(f"audio_{i + 1}" for i in range(MAX_AUDIOS)),
        "image_count", "video_count", "audio_count",
    )
    FUNCTION = "load"
    CATEGORY = "云服务/多路加载"
    OUTPUT_NODE = False

    def load(self, image_paths="", video_paths="", audio_paths="", auto_classify=True):
        # 1. 拆分三组字符串
        raw = {
            "image": _split_paths(image_paths),
            "video": _split_paths(video_paths),
            "audio": _split_paths(audio_paths),
        }

        # 2. 解析真实路径 + 类型分桶
        bucket = {"image": [], "video": [], "audio": []}
        for declared_kind, paths in raw.items():
            for p in paths:
                resolved = _resolve(p)
                if resolved is None:
                    print(f"[MultiFileLoader] 找不到文件，跳过：{p}")
                    continue
                kind = declared_kind
                if auto_classify:
                    guessed = _guess_kind(resolved)
                    if guessed is not None and guessed != declared_kind:
                        print(f"[MultiFileLoader] 自动纠正类型 {declared_kind}→{guessed}: {resolved}")
                        kind = guessed
                bucket[kind].append(resolved)

        # 3. 数量校验（超出直接截断，给个 warning）
        if len(bucket["image"]) > MAX_IMAGES:
            print(f"[MultiFileLoader] 图片超过上限 {MAX_IMAGES}，多余的已忽略")
            bucket["image"] = bucket["image"][:MAX_IMAGES]
        if len(bucket["video"]) > MAX_VIDEOS:
            print(f"[MultiFileLoader] 视频超过上限 {MAX_VIDEOS}，多余的已忽略")
            bucket["video"] = bucket["video"][:MAX_VIDEOS]
        if len(bucket["audio"]) > MAX_AUDIOS:
            print(f"[MultiFileLoader] 音频超过上限 {MAX_AUDIOS}，多余的已忽略")
            bucket["audio"] = bucket["audio"][:MAX_AUDIOS]

        # 4. 逐个加载，未填满的端口填 None
        image_outs = [None] * MAX_IMAGES
        for i, p in enumerate(bucket["image"]):
            try:
                image_outs[i] = _load_image(p)
            except Exception as e:
                print(f"[MultiFileLoader] 加载图片失败 {p}: {e}")

        video_frame_outs = [None] * MAX_VIDEOS
        video_audio_outs = [None] * MAX_VIDEOS
        for i, p in enumerate(bucket["video"]):
            try:
                frames, audio = _load_video(p)
                video_frame_outs[i] = frames
                video_audio_outs[i] = audio
            except Exception as e:
                print(f"[MultiFileLoader] 加载视频失败 {p}: {e}")

        audio_outs = [None] * MAX_AUDIOS
        for i, p in enumerate(bucket["audio"]):
            try:
                audio_outs[i] = _load_audio(p)
            except Exception as e:
                print(f"[MultiFileLoader] 加载音频失败 {p}: {e}")

        # 5. 按 RETURN_NAMES 顺序拼输出
        result: list = list(image_outs)
        for f, a in zip(video_frame_outs, video_audio_outs):
            result.append(f)
            result.append(a)
        result.extend(audio_outs)
        result.extend([
            sum(1 for x in image_outs if x is not None),
            sum(1 for x in video_frame_outs if x is not None),
            sum(1 for x in audio_outs if x is not None),
        ])
        return tuple(result)


NODE_CLASS_MAPPINGS = {
    "MultiFileLoader": MultiFileLoader,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MultiFileLoader": "📦 多路径文件批量加载",
}