# ComfyUI/custom_nodes/zip_output_to_hf.py
#
# Node: ZipOutputToHuggingFace
#
# - Zips the ComfyUI output folder
# - Uploads the zip to a Hugging Face repo
# - Returns a direct download URL as STRING

import os
import time
import zipfile
import tempfile

import folder_paths  # ComfyUI's helper for paths

try:
    from huggingface_hub import HfApi
except ImportError:
    HfApi = None


class ZipOutputToHuggingFace:
    """
    Zip the ComfyUI output folder and upload the archive to Hugging Face.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # Where to upload; default to your repo
                "hf_repo_id": (
                    "STRING",
                    {"default": "saliacoel/MyCustomNodes"},
                ),
            },
            "optional": {
                # Hugging Face token; can be left empty if you use env vars / hf auth login
                "hf_token": (
                    "STRING",
                    {
                        "default": "",
                        "multiline": False,
                    },
                ),
                # Prefix for the zip filename
                "zip_name_prefix": (
                    "STRING",
                    {
                        "default": "comfy_output",
                        "multiline": False,
                    },
                ),
                # Subfolder in the repo to put the zip into
                "remote_dir": (
                    "STRING",
                    {
                        "default": "exports",
                        "multiline": False,
                    },
                ),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("download_url",)
    FUNCTION = "zip_and_upload"
    CATEGORY = "utils/huggingface"
    OUTPUT_NODE = True  # can act as a terminal node in your graph

    def zip_and_upload(
        self,
        hf_repo_id: str,
        hf_token: str = "",
        zip_name_prefix: str = "comfy_output",
        remote_dir: str = "exports",
    ):
        # 1. Check huggingface_hub availability
        if HfApi is None:
            raise RuntimeError(
                "huggingface_hub is not installed.\n"
                "Install it in your ComfyUI environment, for example:\n"
                "  pip install huggingface_hub"
            )

        # 2. Get ComfyUI output directory
        output_dir = folder_paths.get_output_directory()
        if not os.path.isdir(output_dir):
            raise RuntimeError(f"Output directory does not exist: {output_dir}")

        # 3. Ensure there is at least one file to zip
        has_files = False
        for _, _, files in os.walk(output_dir):
            if files:
                has_files = True
                break
        if not has_files:
            raise RuntimeError(
                f"No files found in output directory: {output_dir}"
            )

        # 4. Create zip in a temp directory (so we don't zip the zip)
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        safe_prefix = (zip_name_prefix or "").strip() or "comfy_output"
        zip_filename = f"{safe_prefix}_{timestamp}.zip"

        tmp_dir = tempfile.gettempdir()
        zip_path = os.path.join(tmp_dir, zip_filename)
        zip_abs = os.path.abspath(zip_path)

        # 5. Build the zip archive of everything under output/
        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
            for root, _, files in os.walk(output_dir):
                for fname in files:
                    file_path = os.path.join(root, fname)
                    # Just in case tempdir == output_dir, avoid zipping ourselves
                    if os.path.abspath(file_path) == zip_abs:
                        continue
                    rel_path = os.path.relpath(file_path, output_dir)
                    zipf.write(file_path, arcname=rel_path)

        # 6. Prepare Hugging Face upload
        repo_id = (hf_repo_id or "").strip()
        if not repo_id:
            # Clean up zip if we created it but can't use it
            try:
                os.remove(zip_path)
            except OSError:
                pass
            raise RuntimeError("hf_repo_id cannot be empty")

        remote_dir = (remote_dir or "").strip().strip("/")
        if remote_dir:
            path_in_repo = f"{remote_dir}/{zip_filename}"
        else:
            path_in_repo = zip_filename

        api = HfApi()

        # 7. Upload to Hugging Face
        try:
            api.upload_file(
                path_or_fileobj=zip_path,
                path_in_repo=path_in_repo,
                repo_id=repo_id,
                # token is optional; if empty, hf_hub will use env/CLI token if present
                token=hf_token or None,
            )
        except Exception as e:
            # Clean up local zip on failure
            try:
                os.remove(zip_path)
            except OSError:
                pass
            raise RuntimeError(f"Failed to upload to Hugging Face: {e}")

        # 8. Build a direct download URL
        # This gives the raw file: https://huggingface.co/{repo_id}/resolve/main/{path_in_repo}
        download_url = f"https://huggingface.co/{repo_id}/resolve/main/{path_in_repo}"

        # 9. Clean up temp zip (we only keep the file on HF)
        try:
            os.remove(zip_path)
        except OSError:
            pass

        # Return as a single STRING output
        return (download_url,)


# Register node with ComfyUI
NODE_CLASS_MAPPINGS = {
    "ZipOutputToHuggingFace": ZipOutputToHuggingFace,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ZipOutputToHuggingFace": "Zip Output → Hugging Face",
}