"""Format-aware saving for the unified o1key image workflow.

The generator keeps provider bytes in memory or in ComfyUI's temp directory.
Only the save node promotes those results into a permanent save directory.
"""

from __future__ import annotations

import json
import os
import re
import struct
import threading
import uuid
import zlib
from io import BytesIO
from typing import Any, Iterable

import numpy as np
from PIL import Image
from PIL.PngImagePlugin import PngInfo


SAVE_FORMAT_OPTIONS = ("原始", "png", "jpg", "webp")
SAVE_NAMING_RULE_OPTIONS = ("和主图一致", "自然数字", "自定义前缀")
DEFAULT_SAVE_NAMING_RULE = "自定义前缀"
_FORMAT_EXTENSIONS = {"PNG": "png", "JPEG": "jpg", "WEBP": "webp"}
_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
_SAVE_LOCK = threading.Lock()
_UNSAFE_FILENAME_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]')
_WINDOWS_RESERVED_STEMS = {
    "CON", "PRN", "AUX", "NUL",
    *(f"COM{index}" for index in range(1, 10)),
    *(f"LPT{index}" for index in range(1, 10)),
}


def normalize_save_format(value: Any) -> str:
    normalized = str(value or "原始").strip().lower()
    if normalized == "jpeg":
        normalized = "jpg"
    if normalized == "原始":
        return "原始"
    if normalized not in {"png", "jpg", "webp"}:
        raise ValueError("保存格式仅支持：原始、png、jpg、webp")
    return normalized


def normalize_save_location(value: Any) -> str:
    """Return an output-relative directory or an explicit absolute directory."""

    raw = str(value or "").strip().replace("\\", "/")
    if not raw or raw.lower() == "output" or raw == ".":
        return ""
    drive, _ = os.path.splitdrive(raw)
    if os.path.isabs(raw):
        return os.path.normpath(raw)
    if drive:
        raise ValueError("盘符路径必须填写完整绝对路径，例如 D:/图片")
    parts = [part for part in raw.split("/") if part not in {"", "."}]
    if not parts or any(part == ".." for part in parts):
        raise ValueError("相对保存位置必须是 output 下的子目录，不能包含 ..")
    return os.path.join(*parts)


def _external_preview_descriptor(
    encoded: bytes,
    filename: str,
    folder_paths_module,
) -> dict[str, Any]:
    """Keep external saves previewable without exposing local paths to workflows."""

    temp_root = os.path.abspath(folder_paths_module.get_temp_directory())
    preview_subfolder = os.path.join("o1key_external_preview", uuid.uuid4().hex)
    preview_directory = os.path.join(temp_root, preview_subfolder)
    os.makedirs(preview_directory, exist_ok=False)
    _atomic_write(os.path.join(preview_directory, filename), encoded)
    return {
        "filename": filename,
        "subfolder": preview_subfolder,
        "type": "temp",
        "external_saved": True,
    }


def normalize_naming_rule(value: Any) -> str:
    normalized = str(value or DEFAULT_SAVE_NAMING_RULE).strip()
    if normalized not in SAVE_NAMING_RULE_OPTIONS:
        raise ValueError("命名规则仅支持：和主图一致、自然数字、自定义前缀")
    return normalized


def _safe_filename_stem(value: Any) -> str:
    filename = os.path.basename(str(value or "").strip().replace("\\", "/"))
    stem, _ = os.path.splitext(filename)
    stem = _UNSAFE_FILENAME_CHARS.sub("_", stem).strip(" .")
    stem = stem[:240] or "o1key"
    if stem.upper() in _WINDOWS_RESERVED_STEMS:
        stem = f"_{stem}"
    return stem


def detect_image_format(data: bytes) -> str | None:
    if data.startswith(_PNG_SIGNATURE):
        return "PNG"
    if data.startswith(b"\xff\xd8\xff"):
        return "JPEG"
    if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return "WEBP"
    return None


def _metadata_enabled() -> bool:
    try:
        from comfy.cli_args import args

        return not bool(args.disable_metadata)
    except Exception:
        return True


def _metadata_items(
    prompt: Any = None,
    extra_pnginfo: dict[str, Any] | None = None,
) -> list[tuple[str, str]]:
    if not _metadata_enabled():
        return []
    items: list[tuple[str, str]] = []
    if prompt is not None:
        items.append(("prompt", json.dumps(prompt)))
    if isinstance(extra_pnginfo, dict):
        for key, value in extra_pnginfo.items():
            items.append((str(key), json.dumps(value)))
    return items


def _has_workflow_metadata(extra_pnginfo: dict[str, Any] | None) -> bool:
    """Match ComfyUI's recoverable image contract.

    The frontend metadata parser restores workflows from PNG and WebP, but not
    from JPEG. Native SaveImage always writes PNG, so metadata-bearing JPEG
    results must use that same container instead of an EXIF-only JPEG.
    """

    return (
        _metadata_enabled()
        and isinstance(extra_pnginfo, dict)
        and extra_pnginfo.get("workflow") is not None
    )


def _png_text_chunk(key: str, value: str) -> bytes:
    payload = key.encode("latin-1", "replace") + b"\0" + value.encode(
        "latin-1", "replace"
    )
    chunk_type = b"tEXt"
    return (
        struct.pack(">I", len(payload))
        + chunk_type
        + payload
        + struct.pack(">I", zlib.crc32(chunk_type + payload) & 0xFFFFFFFF)
    )


def inject_png_metadata(data: bytes, metadata: Iterable[tuple[str, str]]) -> bytes:
    """Append tEXt chunks without decoding or recompressing PNG pixels."""
    chunks = [_png_text_chunk(key, value) for key, value in metadata]
    if not chunks or not data.startswith(_PNG_SIGNATURE):
        return data

    offset = len(_PNG_SIGNATURE)
    while offset + 12 <= len(data):
        length = struct.unpack(">I", data[offset : offset + 4])[0]
        end = offset + 12 + length
        if end > len(data):
            return data
        if data[offset + 4 : offset + 8] == b"IEND":
            return data[:offset] + b"".join(chunks) + data[offset:]
        offset = end
    return data


def _create_exif_bytes(metadata: Iterable[tuple[str, str]]) -> bytes:
    items = list(metadata)
    if not items:
        return b""
    exif = Image.Exif()
    extra_index = 0
    for key, value in items:
        if key == "prompt":
            exif[0x0110] = f"prompt:{value}"
        else:
            exif[0x010F - extra_index] = f"{key}:{value}"
            extra_index += 1
    try:
        return exif.tobytes()
    except (OSError, OverflowError, ValueError):
        # JPEG APP1 and some Pillow EXIF writers have practical size limits.
        # The image must still save; PNG remains the fully lossless metadata path.
        return b""


def inject_jpeg_metadata(data: bytes, metadata: Iterable[tuple[str, str]]) -> bytes:
    """Insert an EXIF APP1 segment without recompressing JPEG pixels."""
    if not data.startswith(b"\xff\xd8"):
        return data
    exif = _create_exif_bytes(metadata)
    if not exif or len(exif) + 2 > 0xFFFF:
        return data
    segment = b"\xff\xe1" + struct.pack(">H", len(exif) + 2) + exif
    return data[:2] + segment + data[2:]


def _webp_chunks(data: bytes) -> list[tuple[bytes, bytes]] | None:
    if len(data) < 12 or data[:4] != b"RIFF" or data[8:12] != b"WEBP":
        return None
    chunks: list[tuple[bytes, bytes]] = []
    offset = 12
    while offset + 8 <= len(data):
        kind = data[offset : offset + 4]
        length = struct.unpack("<I", data[offset + 4 : offset + 8])[0]
        start = offset + 8
        end = start + length
        if end > len(data):
            return None
        chunks.append((kind, data[start:end]))
        offset = end + (length & 1)
    return chunks


def _pack_webp_chunk(kind: bytes, payload: bytes) -> bytes:
    padding = b"\0" if len(payload) & 1 else b""
    return kind + struct.pack("<I", len(payload)) + payload + padding


def inject_webp_metadata(data: bytes, metadata: Iterable[tuple[str, str]]) -> bytes:
    """Add/replace a WebP EXIF chunk without recompressing image data."""
    exif = _create_exif_bytes(metadata)
    chunks = _webp_chunks(data)
    if not exif or chunks is None:
        return data

    chunks = [(kind, payload) for kind, payload in chunks if kind != b"EXIF"]
    vp8x_index = next((i for i, item in enumerate(chunks) if item[0] == b"VP8X"), None)
    if vp8x_index is None:
        try:
            with Image.open(BytesIO(data)) as image:
                width, height = image.size
                has_alpha = "A" in image.getbands()
        except Exception:
            return data
        flags = 0x08 | (0x10 if has_alpha else 0)
        vp8x = bytes((flags, 0, 0, 0)) + (width - 1).to_bytes(3, "little") + (
            height - 1
        ).to_bytes(3, "little")
        chunks.insert(0, (b"VP8X", vp8x))
    else:
        kind, payload = chunks[vp8x_index]
        if len(payload) != 10:
            return data
        chunks[vp8x_index] = (kind, bytes((payload[0] | 0x08,)) + payload[1:])

    chunks.append((b"EXIF", exif))
    body = b"WEBP" + b"".join(_pack_webp_chunk(kind, payload) for kind, payload in chunks)
    return b"RIFF" + struct.pack("<I", len(body)) + body


def inject_workflow_metadata(
    data: bytes,
    image_format: str,
    prompt: Any = None,
    extra_pnginfo: dict[str, Any] | None = None,
) -> bytes:
    metadata = _metadata_items(prompt, extra_pnginfo)
    if not metadata:
        return data
    if image_format == "PNG":
        return inject_png_metadata(data, metadata)
    if image_format == "JPEG":
        return inject_jpeg_metadata(data, metadata)
    if image_format == "WEBP":
        return inject_webp_metadata(data, metadata)
    return data


def _atomic_write(path: str, data: bytes) -> None:
    temporary = f"{path}.{uuid.uuid4().hex}.tmp"
    reserved = False
    try:
        descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL)
        os.close(descriptor)
        reserved = True
        with open(temporary, "wb") as handle:
            handle.write(data)
        os.replace(temporary, path)
        reserved = False
    finally:
        if os.path.exists(temporary):
            try:
                os.remove(temporary)
            except OSError:
                pass
        if reserved and os.path.exists(path):
            try:
                os.remove(path)
            except OSError:
                pass


def _tensor_uint8(image_tensor) -> np.ndarray:
    return np.clip(255.0 * image_tensor.detach().cpu().numpy(), 0, 255).astype(np.uint8)


def _raw_bytes_match_tensor(data: bytes, image_tensor) -> bool:
    try:
        array = _tensor_uint8(image_tensor)
        mode = "RGBA" if array.shape[-1] == 4 else "RGB"
        with Image.open(BytesIO(data)) as source:
            source.load()
            decoded = np.asarray(source.convert(mode), dtype=np.uint8)
        return decoded.shape == array.shape and np.array_equal(decoded, array)
    except Exception:
        return False


def _source_for_tensor(images, index: int) -> tuple[bytes | None, str | None, str | None]:
    metadata = getattr(images, "_o1key_source_metadata", None)
    if not isinstance(metadata, list) or index >= len(metadata):
        return None, None, None
    item = metadata[index]
    if not isinstance(item, dict):
        return None, None, None
    source_filename = item.get("filename")
    if not isinstance(source_filename, str) or not source_filename.strip():
        source_filename = None
    if item.get("modified"):
        return None, None, source_filename

    data = item.get("bytes")
    if not isinstance(data, bytes):
        source_path = item.get("path")
        if isinstance(source_path, str) and os.path.isfile(source_path):
            try:
                with open(source_path, "rb") as handle:
                    data = handle.read()
            except OSError:
                data = None
    if not isinstance(data, bytes) or not _raw_bytes_match_tensor(data, images[index]):
        data = None

    image_format = detect_image_format(data) if data else None
    if image_format is None:
        hinted = str(item.get("format") or "").upper()
        if hinted == "JPG":
            hinted = "JPEG"
        image_format = hinted if hinted in _FORMAT_EXTENSIONS else None
    return data, image_format, source_filename


def _pil_from_tensor(image_tensor) -> Image.Image:
    return Image.fromarray(_tensor_uint8(image_tensor))


def _encode_image(
    image: Image.Image,
    image_format: str,
    prompt: Any = None,
    extra_pnginfo: dict[str, Any] | None = None,
) -> bytes:
    output = BytesIO()
    metadata = _metadata_items(prompt, extra_pnginfo)
    working = image
    converted = None
    try:
        if image_format == "PNG":
            pnginfo = None
            if metadata:
                pnginfo = PngInfo()
                for key, value in metadata:
                    pnginfo.add_text(key, value)
            working.save(output, format="PNG", pnginfo=pnginfo, compress_level=4)
        elif image_format == "JPEG":
            if working.mode != "RGB":
                converted = working.convert("RGB")
                working = converted
            kwargs: dict[str, Any] = {
                "format": "JPEG",
                "quality": 100,
                "subsampling": 0,
                "optimize": True,
            }
            exif = _create_exif_bytes(metadata)
            if exif:
                kwargs["exif"] = exif
            working.save(output, **kwargs)
        elif image_format == "WEBP":
            kwargs = {"format": "WEBP", "lossless": True, "quality": 100, "method": 4}
            exif = _create_exif_bytes(metadata)
            if exif:
                kwargs["exif"] = exif
            working.save(output, **kwargs)
        else:
            raise ValueError(f"不支持的图像格式：{image_format}")
        return output.getvalue()
    finally:
        if converted is not None:
            converted.close()


def _save_sources(
    sources: list[tuple[bytes | None, str | None, Image.Image, str | None]],
    filename_prefix: str,
    save_format: str,
    output_directory: str,
    folder_paths_module,
    prompt: Any = None,
    extra_pnginfo: dict[str, Any] | None = None,
    save_location: str = "",
    naming_rule: str = DEFAULT_SAVE_NAMING_RULE,
    main_filename: str | None = None,
) -> list[dict[str, Any]]:
    if not sources:
        return []
    width, height = sources[0][2].size
    normalized_location = normalize_save_location(save_location)
    external_location = bool(normalized_location and os.path.isabs(normalized_location))
    relative_location = "" if external_location else normalized_location
    save_root = normalized_location if external_location else output_directory
    selected_rule = normalize_naming_rule(naming_rule)
    filename = ""
    counter = 1
    if selected_rule == "自定义前缀":
        effective_prefix = filename_prefix or "o1key"
        if relative_location:
            effective_prefix = os.path.join(relative_location, effective_prefix)
        full_output_folder, filename, counter, subfolder, _ = (
            folder_paths_module.get_save_image_path(
                effective_prefix,
                save_root,
                width,
                height,
            )
        )
    else:
        full_output_folder = os.path.join(save_root, relative_location)
        subfolder = relative_location
    os.makedirs(full_output_folder, exist_ok=True)
    selected = normalize_save_format(save_format)
    main_stem = _safe_filename_stem(
        main_filename or next((item[3] for item in sources if item[3]), None)
    )
    results: list[dict[str, Any]] = []
    natural_number = 1
    main_number = 0
    for batch_number, (raw, source_format, image, _source_filename) in enumerate(sources):
        target_format = source_format if selected == "原始" else {
            "png": "PNG",
            "jpg": "JPEG",
            "webp": "WEBP",
        }[selected]
        if target_format not in _FORMAT_EXTENSIONS:
            target_format = "PNG"
        if target_format == "JPEG" and _has_workflow_metadata(extra_pnginfo):
            target_format = "PNG"
        extension = _FORMAT_EXTENSIONS[target_format]
        if selected == "原始" and raw is not None and detect_image_format(raw) == target_format:
            encoded = inject_workflow_metadata(
                raw,
                target_format,
                prompt=prompt,
                extra_pnginfo=extra_pnginfo,
            )
        else:
            encoded = _encode_image(
                image,
                target_format,
                prompt=prompt,
                extra_pnginfo=extra_pnginfo,
            )

        with _SAVE_LOCK:
            while True:
                if selected_rule == "自然数字":
                    saved_name = f"{natural_number}.{extension}"
                    natural_number += 1
                elif selected_rule == "和主图一致":
                    suffix = "" if main_number == 0 else str(main_number)
                    saved_name = f"{main_stem}{suffix}.{extension}"
                    main_number += 1
                else:
                    filename_with_batch_num = filename.replace(
                        "%batch_num%", str(batch_number)
                    )
                    saved_name = f"{filename_with_batch_num}_{counter:05}_.{extension}"
                    counter += 1
                output_path = os.path.join(full_output_folder, saved_name)
                if os.path.exists(output_path):
                    continue
                try:
                    _atomic_write(output_path, encoded)
                    break
                except FileExistsError:
                    continue
        if external_location:
            results.append(
                _external_preview_descriptor(encoded, saved_name, folder_paths_module)
            )
        else:
            results.append({"filename": saved_name, "subfolder": subfolder, "type": "output"})
    return results


def save_tensor_images(
    images,
    filename_prefix: str,
    save_format: str,
    output_directory: str,
    folder_paths_module,
    prompt: Any = None,
    extra_pnginfo: dict[str, Any] | None = None,
    save_location: str = "",
    naming_rule: str = DEFAULT_SAVE_NAMING_RULE,
    main_filename: str | None = None,
) -> list[dict[str, Any]]:
    sources: list[tuple[bytes | None, str | None, Image.Image, str | None]] = []
    try:
        for index, image_tensor in enumerate(images):
            raw, source_format, source_filename = _source_for_tensor(images, index)
            sources.append((raw, source_format, _pil_from_tensor(image_tensor), source_filename))
        return _save_sources(
            sources,
            filename_prefix,
            save_format,
            output_directory,
            folder_paths_module,
            prompt,
            extra_pnginfo,
            save_location,
            naming_rule,
            main_filename or getattr(images, "_o1key_main_filename", None),
        )
    finally:
        for _, _, image, _ in sources:
            image.close()


def save_temp_images(
    source_paths: list[str],
    filename_prefix: str,
    save_format: str,
    output_directory: str,
    folder_paths_module,
    prompt: Any = None,
    extra_pnginfo: dict[str, Any] | None = None,
    save_location: str = "",
    naming_rule: str = DEFAULT_SAVE_NAMING_RULE,
    main_filename: str | None = None,
) -> list[dict[str, Any]]:
    sources: list[tuple[bytes | None, str | None, Image.Image, str | None]] = []
    try:
        for source_path in source_paths:
            with open(source_path, "rb") as handle:
                raw = handle.read()
            image_format = detect_image_format(raw)
            if image_format not in _FORMAT_EXTENSIONS:
                raise ValueError(f"临时结果不是受支持的图像：{os.path.basename(source_path)}")
            with Image.open(BytesIO(raw)) as opened:
                opened.load()
                image = opened.copy()
            sources.append((raw, image_format, image, None))
        return _save_sources(
            sources,
            filename_prefix,
            save_format,
            output_directory,
            folder_paths_module,
            prompt,
            extra_pnginfo,
            save_location,
            naming_rule,
            main_filename,
        )
    finally:
        for _, _, image, _ in sources:
            image.close()


__all__ = [
    "SAVE_FORMAT_OPTIONS",
    "SAVE_NAMING_RULE_OPTIONS",
    "detect_image_format",
    "inject_workflow_metadata",
    "normalize_save_format",
    "normalize_save_location",
    "normalize_naming_rule",
    "save_temp_images",
    "save_tensor_images",
]