"""
ComfyUI model & node registry.

Populate this file via the Admin UI (Models / Custom Nodes / Python Deps tabs)
and click Deploy. You can also edit it by hand — the file is loaded at deploy
time on each runner.

SECTIONS
  1. Core config           — ComfyUI commit pin
  2. Custom nodes          — list of [repo_url, commit_hash]
  3. Models                — dict of {folder_name: [[filename, url], ...]}
  4. Custom node models    — models that live inside `custom_nodes/<node>/...`
  5. Optional S3 override  — load all of the above from S3 at runtime
"""

import json
import os

# =============================================================================
# 1. CORE CONFIG
# =============================================================================

COMFY_COMMIT_HASH = "acd718598eca0b944a1a7a82072a9dec40d3d4f7"  # ComfyUI v0.19.0
REPO_NAME = f"comfy-{COMFY_COMMIT_HASH}-1"

# =============================================================================
# 2. CUSTOM NODES
# =============================================================================
# Format: [repo_url, commit_hash]
# Private repos: use https://USERNAME:TOKEN@github.com/...
# Example:
#   ["https://github.com/cubiq/ComfyUI_essentials", "9d9f4bedfc9f0321c19faf71855e228c93bd0dc9"],

CUSTOM_NODES: list[list[str]] = [
]

# =============================================================================
# 3. MODELS
# =============================================================================
# Format per folder: [filename_on_disk, download_url]
# The file is placed at: <comfy>/models/<folder>/<filename_on_disk>
# Sub-folders are allowed in the filename, e.g. "klein/my-lora.safetensors".

MODELS: dict[str, list[list[str]]] = {
    "diffusion_models": [
    ],
    "checkpoints": [
    ],
    "vae": [
    ],
    "text_encoders": [
    ],
    "loras": [
    ],
    "clip_vision": [
    ],
    "controlnet": [
    ],
    "unet": [
    ],
    "ipadapter": [
    ],
    "upscale_models": [
    ],
}

# =============================================================================
# 4. CUSTOM NODE MODELS
# =============================================================================
# Models that go inside custom_nodes/<node_name>/models/

CUSTOM_NODE_MODELS: dict[str, list[list[str]]] = {
}

# =============================================================================
# 5. S3 CONFIG LOADER (optional)
# =============================================================================
# If S3_CONFIG_BUCKET is set on the runner, this fetches the registry from a
# JSON file in S3 instead of using the lists above. Useful for managing the
# registry centrally across multiple deployments.

def load_config_from_s3():
    """Override MODELS / CUSTOM_NODES from S3 config.json if S3_CONFIG_BUCKET is set."""
    global MODELS, CUSTOM_NODES, CUSTOM_NODE_MODELS

    bucket = os.environ.get("S3_CONFIG_BUCKET")
    if not bucket:
        return

    import boto3

    key = os.environ.get("S3_CONFIG_KEY", "comfyui/config.json")
    region = os.environ.get("S3_CONFIG_REGION", "us-east-1")
    endpoint_url = os.environ.get("S3_ENDPOINT_URL")

    s3_kwargs = {"region_name": region}
    if endpoint_url:
        s3_kwargs["endpoint_url"] = endpoint_url

    s3 = boto3.client("s3", **s3_kwargs)

    print(f"Loading ComfyUI config from s3://{bucket}/{key}")
    response = s3.get_object(Bucket=bucket, Key=key)
    config = json.loads(response["Body"].read().decode("utf-8"))

    if "models" in config:
        MODELS = config["models"]
    if "custom_nodes" in config:
        CUSTOM_NODES = config["custom_nodes"]
    if "custom_node_models" in config:
        CUSTOM_NODE_MODELS = config["custom_node_models"]


load_config_from_s3()
