Azornes-Comfyui-Model-Resolver/core/settings.py

"""Shared settings helpers for Model Resolver."""

from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any, Dict, Iterable, Mapping, Optional
from .path_utils import write_json_atomic, read_json_safe


SETTINGS_FILE = Path(__file__).resolve().parents[1] / "model_resolver_settings.json"

from .type_utils import to_bool, CATEGORY_MAP, normalize_download_category

DOWNLOAD_PATH_MODES = {"suggested", "template", "manual"}
DOWNLOAD_BACKENDS = {"python", "aria2"}

DEFAULT_DOWNLOAD_PATH_TEMPLATES: Dict[str, str] = {
    "loras": "{base_model}/{first_tag}",
    "checkpoints": "{base_model}",
    "embeddings": "{base_model}",
    "diffusion_models": "{base_model}",
    "text_encoders": "",
    "controlnet": "{base_model}",
    "vae": "",
    "upscale_models": "",
    "sams": "",
    "ultralytics": "",
}

DEFAULT_ROOT_KEYS = {
    "loras": "default_lora_root",
    "checkpoints": "default_checkpoint_root",
    "diffusion_models": "default_unet_root",
    "embeddings": "default_embedding_root",
    "text_encoders": "default_text_encoder_root",
    "vae": "default_vae_root",
    "upscale_models": "default_upscale_model_root",
}

TEMPLATE_KEY_ALIASES = {
    "loras": ("loras", "lora"),
    "checkpoints": ("checkpoints", "checkpoint"),
    "embeddings": ("embeddings", "embedding"),
    "diffusion_models": ("diffusion_models", "diffusion_model", "unet"),
    "text_encoders": ("text_encoders", "text_encoder", "clip", "clips", "clip_gguf"),
    "controlnet": ("controlnet", "control_net"),
    "vae": ("vae",),
    "upscale_models": ("upscale_models", "upscale_model", "upscaler"),
    "sams": ("sams", "sam", "sam_model", "sam_models"),
    "ultralytics": ("ultralytics", "ultralytics_bbox", "ultralytics_segm", "yolo"),
}

PRIORITY_TAGS = (
    "concept",
    "style",
    "character",
    "clothing",
    "pose",
    "object",
    "vehicle",
    "artist",
    "celebrity",
)

_INVALID_SEGMENT_CHARS_RE = re.compile(r'[<>:"|?*\x00-\x1f]+')
_UNKNOWN_PLACEHOLDER_RE = re.compile(r"\{[^{}]+\}")


def bool_setting(value: Any, default: bool = True) -> bool:
    return to_bool(value, default)


def sanitize_folder_name(value: Any, fallback: str = "") -> str:
    text = str(value or "").strip()
    if not text:
        text = fallback
    text = text.replace("/", "_").replace("\\", "_")
    text = _INVALID_SEGMENT_CHARS_RE.sub("_", text)
    text = re.sub(r"\s+", " ", text).strip(" .")
    if text in {"", ".", ".."}:
        text = fallback
    return text.strip(" .")


def normalize_relative_subfolder(path_value: Any) -> str:
    text = str(path_value or "").replace("\\", "/").strip()
    if not text:
        return ""
    parts = []
    for raw_part in text.split("/"):
        part = sanitize_folder_name(raw_part)
        if part and part not in {".", ".."}:
            parts.append(part)
    return "/".join(parts)


def normalize_download_path_template(template: Any) -> str:
    text = str(template or "").replace("\\", "/").strip()
    if not text:
        return ""
    parts = []
    for raw_part in text.split("/"):
        part = raw_part.strip()
        if not part or part in {".", ".."}:
            continue
        parts.append(part)
    return "/".join(parts)


def _read_settings_file() -> Dict[str, Any]:
    data = read_json_safe(str(SETTINGS_FILE), {})
    return data if isinstance(data, dict) else {}


def _coerce_dict(value: Any) -> Dict[str, Any]:
    if isinstance(value, dict):
        return dict(value)
    if isinstance(value, str) and value.strip():
        try:
            parsed = json.loads(value)
        except json.JSONDecodeError:
            return {}
        if isinstance(parsed, dict):
            return dict(parsed)
    return {}


def normalize_download_path_templates(value: Any) -> Dict[str, str]:
    raw_templates = _coerce_dict(value)
    normalized = dict(DEFAULT_DOWNLOAD_PATH_TEMPLATES)
    for raw_key, raw_template in raw_templates.items():
        key = normalize_download_category(str(raw_key))
        normalized[key] = normalize_download_path_template(raw_template)
    return normalized


def normalize_base_model_path_mappings(value: Any) -> Dict[str, str]:
    raw_mappings = _coerce_dict(value)
    mappings: Dict[str, str] = {}
    for raw_key, raw_value in raw_mappings.items():
        key = str(raw_key or "").strip()
        value = str(raw_value or "").strip()
        if key and value:
            mappings[key] = value
    return mappings


def normalize_download_path_mode(value: Any, auto_fill_subfolder: Any = None) -> str:
    mode = str(value or "").strip().lower()
    if mode in DOWNLOAD_PATH_MODES:
        return mode
    if auto_fill_subfolder is False:
        return "manual"
    return "suggested"


def normalize_download_backend(value: Any) -> str:
    backend = str(value or "").strip().lower()
    return backend if backend in DOWNLOAD_BACKENDS else "python"


def normalize_settings(settings: Optional[Mapping[str, Any]]) -> Dict[str, Any]:
    data = dict(settings or {})
    data["download_backend"] = normalize_download_backend(data.get("download_backend"))
    data["aria2c_path"] = str(data.get("aria2c_path") or "").strip()
    data["aria2_auto_stop_daemon"] = bool_setting(
        data.get("aria2_auto_stop_daemon"), True
    )
    data["auto_refresh_comfy_models_after_apply"] = bool_setting(
        data.get("auto_refresh_comfy_models_after_apply"), True
    )
    data["download_path_mode"] = normalize_download_path_mode(
        data.get("download_path_mode"), data.get("auto_fill_subfolder")
    )
    data["download_path_templates"] = normalize_download_path_templates(
        data.get("download_path_templates")
    )
    data["base_model_path_mappings"] = normalize_base_model_path_mappings(
        data.get("base_model_path_mappings")
    )
    for key in DEFAULT_ROOT_KEYS.values():
        data[key] = str(data.get(key) or "").strip()
    return data


def load_settings() -> Dict[str, Any]:
    return normalize_settings(_read_settings_file())


def save_settings(payload: Mapping[str, Any]) -> Dict[str, Any]:
    current = _read_settings_file()
    current.update({str(key): value for key, value in payload.items() if key})
    normalized = normalize_settings(current)
    write_json_atomic(str(SETTINGS_FILE), normalized, indent=2)
    return normalized


def _template_for_category(settings: Mapping[str, Any], category: str) -> str:
    templates = normalize_download_path_templates(settings.get("download_path_templates"))
    category_key = normalize_download_category(category)
    aliases = TEMPLATE_KEY_ALIASES.get(category_key, (category_key,))
    for alias in aliases:
        if alias in templates:
            return templates.get(alias, "")
    return templates.get(category_key, "")


def _listify_tags(value: Any) -> Iterable[str]:
    if isinstance(value, (list, tuple, set)):
        for item in value:
            text = str(item or "").strip()
            if text:
                yield text
        return
    if isinstance(value, str):
        for item in re.split(r"[,;]+", value):
            text = item.strip()
            if text:
                yield text


def _normalize_tag(value: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", str(value or "").lower())


def _resolve_base_model_mapping(mappings: Mapping[str, str], base_model: Any) -> str:
    base_text = str(base_model or "")
    if base_text in mappings:
        return mappings[base_text]

    base_token = _normalize_tag(base_text)
    if not base_token:
        return base_text

    normalized_mappings = {
        _normalize_tag(key): value
        for key, value in mappings.items()
        if _normalize_tag(key)
    }
    if base_token in normalized_mappings:
        return normalized_mappings[base_token]

    for key_token, value in sorted(
        normalized_mappings.items(),
        key=lambda item: len(item[0]),
        reverse=True,
    ):
        if len(key_token) < 4:
            continue
        if base_token.startswith(key_token) or key_token in base_token or base_token in key_token:
            return value

    return base_text


def _first_tag(tags: Any) -> str:
    tag_list = list(_listify_tags(tags))
    if not tag_list:
        return "no tags"
    normalized_tags = {tag: _normalize_tag(tag) for tag in tag_list}
    for priority in PRIORITY_TAGS:
        priority_token = _normalize_tag(priority)
        for tag, normalized in normalized_tags.items():
            if normalized == priority_token:
                return tag
    return tag_list[0]


def _creator_name(metadata: Mapping[str, Any]) -> str:
    creator = metadata.get("creator")
    if isinstance(creator, Mapping):
        name = creator.get("username") or creator.get("name")
        if name:
            return str(name)
    if isinstance(creator, str) and creator.strip():
        return creator
    for key in ("author", "username", "creator_username", "creator_name"):
        value = metadata.get(key)
        if value:
            return str(value)
    repo_id = metadata.get("repo_id") or metadata.get("repo")
    if isinstance(repo_id, str) and "/" in repo_id:
        return repo_id.split("/", 1)[0]
    return "Anonymous"


def _filename_stem(metadata: Mapping[str, Any]) -> str:
    filename = str(metadata.get("filename") or metadata.get("file_name") or "").strip()
    if not filename:
        return ""
    return Path(filename).stem


def calculate_template_subfolder(
    category: str,
    metadata: Optional[Mapping[str, Any]] = None,
    settings: Optional[Mapping[str, Any]] = None,
) -> str:
    active_settings = normalize_settings(settings if settings is not None else load_settings())
    template = _template_for_category(active_settings, category)
    if not template:
        return ""

    model_metadata = dict(metadata or {})
    base_model = (
        model_metadata.get("base_model")
        or model_metadata.get("baseModel")
        or model_metadata.get("base_model_context")
        or "Unknown Base Model"
    )
    mappings = normalize_base_model_path_mappings(
        active_settings.get("base_model_path_mappings")
    )
    mapped_base_model = _resolve_base_model_mapping(mappings, base_model)
    model_name = (
        model_metadata.get("model_name")
        or model_metadata.get("name")
        or _filename_stem(model_metadata)
        or "Model"
    )
    version_name = (
        model_metadata.get("version_name")
        or model_metadata.get("versionName")
        or model_metadata.get("version")
        or ""
    )

    replacements = {
        "{base_model}": (
            normalize_relative_subfolder(mapped_base_model)
            or sanitize_folder_name(mapped_base_model, "Unknown Base Model")
        ),
        "{author}": sanitize_folder_name(_creator_name(model_metadata), "Anonymous"),
        "{first_tag}": sanitize_folder_name(
            _first_tag(model_metadata.get("tags")), "no tags"
        ),
        "{model_name}": sanitize_folder_name(model_name, "Model"),
        "{version_name}": sanitize_folder_name(version_name, ""),
    }

    formatted = template
    for token, replacement in replacements.items():
        formatted = formatted.replace(token, replacement)
    formatted = _UNKNOWN_PLACEHOLDER_RE.sub("", formatted)
    return normalize_relative_subfolder(formatted)


def resolve_download_subfolder(
    category: str,
    requested_subfolder: Any = "",
    metadata: Optional[Mapping[str, Any]] = None,
    settings: Optional[Mapping[str, Any]] = None,
) -> str:
    requested = str(requested_subfolder or "").strip()
    if requested:
        return requested

    active_settings = normalize_settings(settings if settings is not None else load_settings())
    if active_settings.get("download_path_mode") != "template":
        return ""

    return calculate_template_subfolder(category, metadata, active_settings)


def get_default_root_for_category(
    category: str,
    settings: Optional[Mapping[str, Any]] = None,
) -> str:
    active_settings = normalize_settings(settings if settings is not None else load_settings())
    key = DEFAULT_ROOT_KEYS.get(normalize_download_category(category))
    if not key:
        return ""
    return str(active_settings.get(key) or "").strip()