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

"""
Core Resolver Module

Integrates all components to provide high-level API for model linking.
"""

import os
import re
import json
from typing import Dict, Any, List, Optional, Tuple, Callable
from urllib.parse import unquote

from .log_system.log_funcs import (
    log_debug,
    log_info,
    log_warn,
    log_error,
    log_exception,
)
from .scanner import get_model_files
from .workflow_analyzer import (
    NODE_TYPE_TO_CATEGORY_HINTS,
    analyze_workflow_models,
    identify_missing_models,
)
from .matcher import find_matches
from .workflow_updater import update_workflow_nodes
from .sources.civitai import resolve_urn

# Regex patterns for URL extraction (matches HuggingFace and CivitAI URLs)
URL_PATTERN = re.compile(r'(https?://(?:huggingface\.co|civitai\.com)[^\s"\'<>\)\\]+)')

# Model file extensions to look for
MODEL_EXTENSIONS = (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx", ".gguf")


def get_path_identity(path: str) -> str:
    """Return a stable path identity across symlinks/junctions."""
    if not path:
        return ""

    try:
        return os.path.normcase(os.path.realpath(os.path.abspath(path)))
    except (OSError, ValueError):
        return os.path.normcase(os.path.abspath(path))


def strip_known_model_extension(filename: str) -> str:
    """Strip only known model extensions, preserving names like v4.0."""
    if not isinstance(filename, str):
        return ""

    lowered = filename.lower()
    for ext in MODEL_EXTENSIONS:
        if lowered.endswith(ext):
            return filename[: -len(ext)]
    return filename


def get_workflow_url_info_for_filename(
    workflow_urls: Dict[str, Dict[str, Any]], filename: str
) -> Optional[Dict[str, Any]]:
    if filename in workflow_urls:
        return workflow_urls[filename]

    filename_stem = strip_known_model_extension(os.path.basename(filename)).lower()
    if not filename_stem:
        return None

    for workflow_filename, url_info in workflow_urls.items():
        workflow_stem = strip_known_model_extension(
            os.path.basename(workflow_filename)
        ).lower()
        if workflow_stem == filename_stem:
            return url_info

    return None


def normalize_workflow_download_url(url: str) -> str:
    """Convert workflow file-page URLs into direct download URLs when possible."""
    if not isinstance(url, str) or not url:
        return url

    # HuggingFace /blob/ pages return HTML. /resolve/ returns the actual file.
    return re.sub(
        r"(https?://huggingface\.co/[^/]+/[^/]+)/blob/([^/]+)/(.+)",
        r"\1/resolve/\2/\3",
        url,
    )


def workflow_url_points_to_file(url: str, filename: str) -> bool:
    """Return true when a URL appears to reference the specific model file."""
    if not url or not filename:
        return False

    try:
        decoded_url = unquote(url)
    except Exception:
        decoded_url = url

    return filename in decoded_url or unquote(filename) in decoded_url


def search_local_matches(
    target_for_matching: str,
    category: Optional[str] = None,
    similarity_threshold: float = 0.0,
    max_matches_per_model: int = 10,
    force_rescan: bool = False,
) -> List[Dict[str, Any]]:
    """
    Search local model files using the same matcher as workflow analysis.

    Args:
        target_for_matching: Filename/path to match against local files
        category: Optional category hint to prioritize/filter candidates
        similarity_threshold: Minimum similarity score (0.0 to 1.0)
        max_matches_per_model: Maximum number of matches to return

    Returns:
        Deduplicated list of local matches sorted by similarity
    """
    available_models = get_model_files(force_rescan=force_rescan)

    candidates = available_models
    if category and category != "unknown":
        candidates = [m for m in available_models if m.get("category") == category]
        candidates.extend(
            [m for m in available_models if m.get("category") != category]
        )

    matches = find_matches(
        target_for_matching,
        candidates,
        threshold=similarity_threshold,
        max_results=max_matches_per_model,
    )

    seen_absolute_paths = {}
    deduplicated_matches = []
    for match in matches:
        model_dict = match["model"]
        absolute_path = model_dict.get("path", "")
        path_identity = get_path_identity(absolute_path) if absolute_path else ""
        dedupe_key = path_identity or os.path.normcase(
            model_dict.get("relative_path", "") or match.get("filename", "")
        )

        if dedupe_key not in seen_absolute_paths:
            seen_absolute_paths[dedupe_key] = match
            deduplicated_matches.append(match)
        else:
            existing_match = seen_absolute_paths[dedupe_key]
            if match["confidence"] > existing_match["confidence"]:
                idx = deduplicated_matches.index(existing_match)
                deduplicated_matches[idx] = match
                seen_absolute_paths[dedupe_key] = match

    return deduplicated_matches


def extract_workflow_urls(workflow_json: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
    """
    Extract model URLs from workflow JSON.

    Sources:
    1. node.properties.models array - contains {name, url, directory}
    2. Regex extraction from workflow JSON string - finds HuggingFace/CivitAI URLs

    Args:
        workflow_json: Complete workflow JSON dictionary

    Returns:
        Dict mapping model filename -> {url, directory, source}
    """
    url_map = {}

    # Convert to string for regex search
    workflow_str = json.dumps(workflow_json)

    # Collect all nodes including from subgraphs
    all_nodes = list(workflow_json.get("nodes", []))
    definitions = workflow_json.get("definitions", {})
    subgraphs = definitions.get("subgraphs", [])
    for subgraph in subgraphs:
        subgraph_nodes = subgraph.get("nodes", [])
        all_nodes.extend(subgraph_nodes)

    # 1. Extract from node.properties.models (authoritative source)
    for node in all_nodes:
        node_type = node.get("type", "")
        properties = node.get("properties", {})
        models_list = properties.get("models", [])

        for model_info in models_list:
            if isinstance(model_info, dict):
                name = model_info.get("name", "")
                url = model_info.get("url", "")
                directory = model_info.get("directory", "")

                if name and name not in url_map:
                    url_map[name] = {
                        "url": normalize_workflow_download_url(url),
                        "model_url": url,
                        "directory": directory,
                        "node_type": node_type,
                        "source": "node_properties",
                    }

    # 2. Extract URLs via regex from workflow JSON
    urls_found = URL_PATTERN.findall(workflow_str)

    # Clean URLs (remove trailing characters that may have been captured)
    cleaned_urls = []
    for url in urls_found:
        url = url.split(")")[0].replace("\\n", "").replace("\n", "").strip()
        if url:
            cleaned_urls.append(url)

    # 3. Extract model filenames via regex
    model_pattern = re.compile(
        r"([\w\-\.%]+\.(?:safetensors|ckpt|pt|pth|bin|onnx|gguf))", re.IGNORECASE
    )
    model_files_raw = model_pattern.findall(workflow_str)

    # Clean and decode filenames
    model_files = set()
    model_name_map = {}  # decoded -> original

    for model in model_files_raw:
        cleaned = model.strip()
        if cleaned and cleaned[0].isalnum():
            try:
                decoded = unquote(cleaned)
            except Exception:
                decoded = cleaned
            model_files.add(decoded)
            model_name_map[decoded] = cleaned

    # 4. Match URLs to model filenames
    for model in model_files:
        # Keep authoritative node.properties URLs only when they point to the file.
        # Some workflows store a repo/model-page URL there and a concrete file URL
        # elsewhere in the JSON; in that case the file URL is the usable download.
        if (
            model in url_map
            and url_map[model].get("url")
            and workflow_url_points_to_file(url_map[model].get("url"), model)
        ):
            continue

        original_name = model_name_map.get(model, model)

        for url in cleaned_urls:
            # Check decoded name in URL
            if model in url:
                if (
                    model not in url_map
                    or not url_map[model].get("url")
                    or not workflow_url_points_to_file(url_map[model].get("url"), model)
                ):
                    url_map[model] = {
                        "url": normalize_workflow_download_url(url),
                        "model_url": url,
                        "directory": url_map.get(model, {}).get("directory", ""),
                        "source": "regex",
                    }
                else:
                    url_map[model]["url"] = normalize_workflow_download_url(url)
                    url_map[model]["model_url"] = url
                    url_map[model]["source"] = "regex"
                break
            # Check original (possibly URL-encoded) name in URL
            if original_name in url:
                if (
                    model not in url_map
                    or not url_map[model].get("url")
                    or not workflow_url_points_to_file(
                        url_map[model].get("url"), original_name
                    )
                ):
                    url_map[model] = {
                        "url": normalize_workflow_download_url(url),
                        "model_url": url,
                        "directory": url_map.get(model, {}).get("directory", ""),
                        "source": "regex",
                    }
                else:
                    url_map[model]["url"] = normalize_workflow_download_url(url)
                    url_map[model]["model_url"] = url
                    url_map[model]["source"] = "regex"
                break
            # Check without extension
            model_base = os.path.splitext(model)[0]
            if model_base in url or unquote(model_base) in url:
                if (
                    model not in url_map
                    or not url_map[model].get("url")
                    or not workflow_url_points_to_file(url_map[model].get("url"), model)
                ):
                    url_map[model] = {
                        "url": normalize_workflow_download_url(url),
                        "model_url": url,
                        "directory": url_map.get(model, {}).get("directory", ""),
                        "source": "regex",
                    }
                else:
                    url_map[model]["url"] = normalize_workflow_download_url(url)
                    url_map[model]["model_url"] = url
                    url_map[model]["source"] = "regex"
                break

    return url_map


def parse_huggingface_url(url: str) -> Tuple[Optional[str], Optional[str]]:
    """
    Extract HuggingFace repo and path from URL.

    Args:
        url: HuggingFace URL

    Returns:
        Tuple of (repo_id, file_path) or (None, None) if not valid
    """
    if not url or "huggingface.co" not in url:
        return None, None

    # Pattern: https://huggingface.co/user/repo/resolve/main/path/to/file.safetensors
    match = re.match(
        r"https?://huggingface\.co/([^/]+/[^/]+)/(?:resolve|blob)/[^/]+/(.+)", url
    )
    if match:
        return match.group(1), match.group(2)

    return None, None


def analyze_and_find_matches(
    workflow_json: Dict[str, Any],
    similarity_threshold: float = 0.0,
    max_matches_per_model: int = 10,
    progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
    force_rescan: bool = False,
) -> Dict[str, Any]:
    """
    Main entry point: analyze workflow and find matches for missing models.

    Args:
        workflow_json: Complete workflow JSON dictionary
        similarity_threshold: Minimum similarity score (0.0 to 1.0) for matches
        max_matches_per_model: Maximum number of matches to return per missing model
        force_rescan: If True, bypass the short-lived local model scan cache

    Returns:
        Dictionary with analysis results:
        {
            'missing_models': [
                {
                    'node_id': node ID,
                    'node_type': node type,
                    'widget_index': widget index,
                    'original_path': original path from workflow,
                    'category': model category,
                    'workflow_url': URL from workflow if found,
                    'workflow_directory': directory from workflow if found,
                    'matches': [
                        {
                            'model': model dict from scanner,
                            'filename': model filename,
                            'similarity': similarity score (0.0-1.0),
                            'confidence': confidence percentage (0-100)
                        },
                        ...
                    ]
                },
                ...
            ],
            'total_missing': count of missing models,
            'total_models_analyzed': count of all models in workflow
        }
    """
    if progress_callback:
        progress_callback(
            {
                "stage": "extracting",
                "message": "Extracting workflow model references...",
                "current": 0,
                "total": 0,
            }
        )

    # Extract URLs from workflow (node.properties.models + regex)
    workflow_urls = extract_workflow_urls(workflow_json)
    log_debug(f"Extracted {len(workflow_urls)} URLs from workflow")

    if progress_callback:
        progress_callback(
            {
                "stage": "scanning",
                "message": "Scanning local model index...",
                "current": 0,
                "total": 0,
            }
        )

    # Analyze workflow to find all model references
    # Get available models
    available_models = get_model_files(force_rescan=force_rescan)
    available_models_by_category = {}
    for model in available_models:
        model_category = model.get("category", "")
        if model_category not in available_models_by_category:
            available_models_by_category[model_category] = []
        available_models_by_category[model_category].append(model)

    ordered_candidates_cache: Dict[str, List[Dict[str, Any]]] = {}

    if progress_callback:
        progress_callback(
            {
                "stage": "analyzing",
                "message": "Analyzing workflow nodes...",
                "current": 0,
                "total": 0,
            }
        )

    # Analyze workflow using the same already-scanned model list
    all_model_refs = analyze_workflow_models(
        workflow_json, available_models=available_models
    )

    if progress_callback:
        progress_callback(
            {
                "stage": "identifying",
                "message": "Identifying missing models...",
                "current": 0,
                "total": len(all_model_refs),
            }
        )

    # Identify missing models
    missing_models = identify_missing_models(all_model_refs, available_models)
    resolved_model_refs = [
        model_ref for model_ref in all_model_refs if model_ref.get("exists", False)
    ]

    # Enrich missing models with workflow URLs
    for missing in missing_models:
        original_path = missing.get("original_path", "")
        filename = os.path.basename(original_path)

        url_info = get_workflow_url_info_for_filename(workflow_urls, filename)
        if url_info:
            missing["workflow_url"] = url_info.get("url", "")
            missing["workflow_model_url"] = url_info.get("model_url", "")
            missing["workflow_directory"] = url_info.get("directory", "")
            missing["url_source"] = url_info.get("source", "")

    # Handle URNs: mark for async resolution by frontend
    # No sync CivitAI calls here - frontend will fetch asynchronously
    for missing in missing_models:
        if missing.get("is_urn"):
            missing["needs_urn_resolve"] = True
            urn = missing.get("urn")
            if urn:
                missing["urn_model_id"] = urn.get("model_id")
                missing["urn_version_id"] = urn.get("version_id")
                missing["urn_type"] = urn.get("type", "")

    total_matching_models = len(missing_models) + len(resolved_model_refs)
    if progress_callback:
        progress_callback(
            {
                "stage": "matching",
                "message": "Matching local models...",
                "current": 0,
                "total": total_matching_models,
            }
        )

    local_match_cache: Dict[Tuple[str, str], List[Dict[str, Any]]] = {}

    def get_match_target(model_ref: Dict[str, Any]) -> Optional[str]:
        target_for_matching = (
            model_ref.get("original_path")
            or model_ref.get("expected_filename")
            or model_ref.get("name")
            or model_ref.get("filename")
            or model_ref.get("full_path")
            or ""
        )

        # For URNs, prefer expected_filename for matching
        if model_ref.get("is_urn") and model_ref.get("expected_filename"):
            return model_ref["expected_filename"]

        if isinstance(target_for_matching, str) and target_for_matching.startswith(
            "urn:air:"
        ):
            # Don't fuzzy-match the full URN string against every local model.
            # For worker-asset style URNs, try using the filename-like suffix after "@";
            # otherwise skip local matching until the frontend resolves the URN async.
            urn_suffix = (
                target_for_matching.split("@", 1)[1]
                if "@" in target_for_matching
                else ""
            )
            urn_suffix_ext = os.path.splitext(urn_suffix)[1].lower()
            if urn_suffix and urn_suffix_ext in MODEL_EXTENSIONS:
                return urn_suffix
            return None

        return target_for_matching

    def get_match_category(model_ref: Dict[str, Any]) -> str:
        category = model_ref.get("category")
        if not category or category == "unknown":
            node_type = model_ref.get("node_type", "")
            category = NODE_TYPE_TO_CATEGORY_HINTS.get(node_type, "unknown")
        return category or "unknown"

    def get_candidates_for_category(category: str) -> List[Dict[str, Any]]:
        if not category or category == "unknown":
            return available_models

        candidates = ordered_candidates_cache.get(category)
        if candidates is None:
            preferred = available_models_by_category.get(category, [])
            others = [
                m for m in available_models if m.get("category") != category
            ]
            candidates = preferred + others
            ordered_candidates_cache[category] = candidates

        return candidates

    def deduplicate_matches(matches: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        seen_absolute_paths = {}
        deduplicated_matches = []
        for match in matches:
            model_dict = match["model"]
            absolute_path = model_dict.get("path", "")
            path_identity = get_path_identity(absolute_path) if absolute_path else ""
            dedupe_key = path_identity or os.path.normcase(
                model_dict.get("relative_path", "") or match.get("filename", "")
            )

            if dedupe_key not in seen_absolute_paths:
                seen_absolute_paths[dedupe_key] = match
                deduplicated_matches.append(match)
            else:
                existing_match = seen_absolute_paths[dedupe_key]
                if match["confidence"] > existing_match["confidence"]:
                    idx = deduplicated_matches.index(existing_match)
                    deduplicated_matches[idx] = match
                    seen_absolute_paths[dedupe_key] = match

        return deduplicated_matches

    def find_local_matches_for_ref(model_ref: Dict[str, Any]) -> List[Dict[str, Any]]:
        target_for_matching = get_match_target(model_ref)
        if not target_for_matching:
            return []

        category = get_match_category(model_ref)
        cache_key = (target_for_matching, category)
        if cache_key in local_match_cache:
            return local_match_cache[cache_key]

        matches = find_matches(
            target_for_matching,
            get_candidates_for_category(category),
            threshold=similarity_threshold,
            max_results=max_matches_per_model,
        )
        deduplicated_matches = deduplicate_matches(matches)
        local_match_cache[cache_key] = deduplicated_matches
        return deduplicated_matches

    # Find matches for each missing model
    missing_with_matches = []
    total_missing = len(missing_models)
    for index, missing in enumerate(missing_models, start=1):
        if progress_callback:
            progress_callback(
                {
                    "stage": "matching",
                    "message": f"Analyzing model {index} of {total_missing}",
                    "current": index,
                    "total": total_matching_models,
                    "model_name": missing.get("name")
                    or missing.get("original_path", ""),
                }
            )

        # Skip LoraManager lorAs that already exist locally (exists=True means no linking needed)
        is_lora_v2 = missing.get("is_lora_v2")
        exists = missing.get("exists")
        name = missing.get("name") or missing.get("original_path", "")
        log_debug(f"Checking {name}: is_lora_v2={is_lora_v2}, exists={exists}")

        if is_lora_v2 and exists:
            log_info(f"Skipping LoraManager lora {name} - already exists locally")
            continue

        missing_with_matches.append({
            **missing,
            "matches": find_local_matches_for_ref(missing),
        })

    resolved_with_matches = []
    total_resolved = len(resolved_model_refs)
    for index, resolved in enumerate(resolved_model_refs, start=1):
        if progress_callback:
            progress_callback(
                {
                    "stage": "matching",
                    "message": f"Analyzing resolved model {index} of {total_resolved}",
                    "current": total_missing + index,
                    "total": total_matching_models,
                    "model_name": resolved.get("name")
                    or resolved.get("original_path", ""),
                }
            )

        resolved_with_matches.append({
            **resolved,
            "matches": find_local_matches_for_ref(resolved),
        })

    result = {
        "missing_models": missing_with_matches,
        "resolved_models": resolved_with_matches,
        "total_resolved": len(resolved_with_matches),
        "total_missing": len(missing_with_matches),
        "total_models_analyzed": len(all_model_refs),
    }

    if progress_callback:
        progress_callback(
            {
                "stage": "completed",
                "message": "Analysis complete",
                "current": total_matching_models,
                "total": total_matching_models,
            }
        )

    return result


def apply_resolution(
    workflow_json: Dict[str, Any], resolutions: List[Dict[str, Any]]
) -> Dict[str, Any]:
    """
    Apply model resolutions to workflow.

    Args:
        workflow_json: Workflow JSON dictionary (will be modified)
        resolutions: List of resolution dictionaries:
            {
                'node_id': node ID,
                'widget_index': widget index,
                'resolved_path': absolute path to resolved model,
                'category': model category (optional),
                'resolved_model': model dict from scanner (optional),
                'nested_key': nested key for dict-type widgets (optional)
            }

    Returns:
        Updated workflow JSON dictionary
    """
    # Prepare mappings for workflow_updater
    mappings = []
    for resolution in resolutions:
        mapping = {
            "node_id": resolution.get("node_id"),
            "widget_index": resolution.get("widget_index"),
            "resolved_path": resolution.get("resolved_path"),
            "category": resolution.get("category"),
            "resolved_model": resolution.get("resolved_model"),
            "subgraph_id": resolution.get(
                "subgraph_id"
            ),  # Include subgraph_id for subgraph nodes
            "is_top_level": resolution.get(
                "is_top_level"
            ),  # True for top-level nodes, False for nodes in subgraph definitions
            "is_lora_v2": resolution.get("is_lora_v2"),  # Flag for LoraManager nodes
            "original_lora_name": resolution.get(
                "original_lora_name"
            ),  # Original lora name for LoraManager replacement
            "nested_key": resolution.get(
                "nested_key"
            ),  # For dict-type widgets (e.g. Power Lora Loader)
        }

        # If resolved_model provided, extract path if needed
        if "resolved_model" in resolution and resolution["resolved_model"]:
            resolved_model = resolution["resolved_model"]
            if "path" in resolved_model and not mapping.get("resolved_path"):
                mapping["resolved_path"] = resolved_model["path"]
            if "base_directory" in resolved_model:
                mapping["base_directory"] = resolved_model["base_directory"]

        mappings.append(mapping)

    # Update workflow
    updated_workflow = update_workflow_nodes(workflow_json, mappings)

    return updated_workflow


def get_resolution_summary(workflow_json: Dict[str, Any]) -> Dict[str, Any]:
    """
    Get summary of missing models and matches without applying resolutions.

    This is a convenience method that calls analyze_and_find_matches with defaults.

    Args:
        workflow_json: Complete workflow JSON dictionary

    Returns:
        Same format as analyze_and_find_matches
    """
    return analyze_and_find_matches(workflow_json)