"""Scan ComfyUI workflow JSON files for required model filenames."""

from __future__ import annotations

import json
import logging
import subprocess
from collections.abc import Iterator
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from comfyui_config import COMFYUI_CONTAINER, container_workflow_path

log = logging.getLogger(__name__)

# class_type -> list of (input_field, category)
LOADER_NODES: dict[str, list[tuple[str, str]]] = {
    "CheckpointLoaderSimple": [("ckpt_name", "checkpoint")],
    "CheckpointLoader": [("ckpt_name", "checkpoint")],
    "unCLIPCheckpointLoader": [("ckpt_name", "checkpoint")],
    "ImageOnlyCheckpointLoader": [("ckpt_name", "checkpoint")],
    "VAELoader": [("vae_name", "vae")],
    "LoraLoader": [("lora_name", "lora")],
    "LoraLoaderModelOnly": [("lora_name", "lora")],
    "UNETLoader": [("unet_name", "unet")],
    "CLIPLoader": [("clip_name", "clip")],
    "DualCLIPLoader": [("clip_name1", "clip"), ("clip_name2", "clip")],
    "ControlNetLoader": [("control_net_name", "controlnet")],
    "UpscaleModelLoader": [("model_name", "upscale")],
    "StyleModelLoader": [("style_model_name", "style_models")],
    "LatentUpscaleModelLoader": [("model_name", "latent_upscale")],
    "LTXVAudioVAELoader": [("ckpt_name", "vae")],
    "LTXAVTextEncoderLoader": [("gemma_model", "text_encoder"), ("ltx_model", "checkpoint")],
    "DiffusersLoader": [("model_path", "diffusion_models")],
}

# Nodes where every .safetensors in widgets_values should be collected
MULTI_WIDGET_LOADERS = frozenset({"LTXAVTextEncoderLoader"})


@dataclass
class ModelRef:
    filename: str
    category: str
    workflows: list[str] = field(default_factory=list)


def _is_model_filename(value: Any) -> bool:
    if not isinstance(value, str) or not value.strip():
        return False
    lower = value.lower()
    return lower.endswith(
        (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx", ".sft")
    )


def _infer_category_from_node(class_type: str, filename: str) -> str:
    ct = class_type.lower()
    fn = filename.lower()
    if "lora" in ct or "lora" in fn:
        return "lora"
    if "upscale" in ct or "upscaler" in fn or "upscale" in fn:
        return "latent_upscale"
    if "vae" in ct:
        return "vae"
    if "clip" in ct or "text_encoder" in ct or "gemma" in fn:
        return "text_encoder"
    if "unet" in ct:
        return "unet"
    if "controlnet" in ct:
        return "controlnet"
    return "checkpoint"


def _read_workflow_json(path: Path, workflows_dir: Path) -> dict[str, Any] | None:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except PermissionError:
        try:
            rel = path.relative_to(workflows_dir.resolve())
        except ValueError:
            log.warning("Workflow outside workflows dir, cannot use docker fallback: %s", path)
            return None
        container_path = container_workflow_path(rel.as_posix())
        try:
            result = subprocess.run(
                ["docker", "exec", COMFYUI_CONTAINER, "cat", container_path],
                capture_output=True,
                text=True,
                timeout=30,
                check=False,
            )
        except (OSError, subprocess.TimeoutExpired) as exc:
            log.warning("Docker fallback failed for workflow %s: %s", path, exc)
            return None
        if result.returncode != 0:
            log.warning(
                "Docker fallback failed for workflow %s: %s",
                path,
                result.stderr.strip(),
            )
            return None
        log.debug("Read workflow via docker exec: %s", path.name)
        try:
            data = json.loads(result.stdout)
        except json.JSONDecodeError as exc:
            log.warning("Invalid JSON from docker fallback for %s: %s", path, exc)
            return None
        return data if isinstance(data, dict) else None
    except (OSError, json.JSONDecodeError) as exc:
        log.warning("Failed to read workflow %s: %s", path, exc)
        return None


def _iter_workflow_nodes(data: dict[str, Any]) -> Iterator[dict[str, Any]]:
    nodes = data.get("nodes")
    if isinstance(nodes, list):
        for node in nodes:
            if isinstance(node, dict):
                yield node

    definitions = data.get("definitions")
    if isinstance(definitions, dict):
        subgraphs = definitions.get("subgraphs")
        if isinstance(subgraphs, list):
            for subgraph in subgraphs:
                if not isinstance(subgraph, dict):
                    continue
                inner = subgraph.get("nodes")
                if isinstance(inner, list):
                    for node in inner:
                        if isinstance(node, dict):
                            yield node


def _node_class_type(node: dict[str, Any]) -> str | None:
    ct = node.get("class_type") or node.get("type")
    if isinstance(ct, str) and not ct.isupper() and ct not in ("INT", "FLOAT", "STRING", "COMBO", "BOOLEAN", "*"):
        return ct
    return None


def _extract_from_widgets(
    class_type: str,
    widgets: list[Any],
) -> list[tuple[str, str]]:
    refs: list[tuple[str, str]] = []
    if class_type in MULTI_WIDGET_LOADERS:
        field_specs = LOADER_NODES.get(class_type, [])
        for idx, value in enumerate(widgets):
            if not _is_model_filename(value):
                continue
            if idx < len(field_specs):
                category = field_specs[idx][1]
            else:
                category = _infer_category_from_node(class_type, str(value))
            refs.append((str(value), category))
        return refs

    if class_type in LOADER_NODES and widgets:
        first = widgets[0]
        if _is_model_filename(first):
            refs.append((str(first), LOADER_NODES[class_type][0][1]))
    return refs


def _extract_generic_widgets(class_type: str, widgets: list[Any]) -> list[tuple[str, str]]:
    refs: list[tuple[str, str]] = []
    for value in widgets:
        if _is_model_filename(value):
            refs.append((str(value), _infer_category_from_node(class_type, str(value))))
    return refs


def _extract_from_node(node: dict[str, Any]) -> list[tuple[str, str]]:
    refs: list[tuple[str, str]] = []
    class_type = _node_class_type(node)
    if not class_type:
        return refs

    widgets = node.get("widgets_values")
    if isinstance(widgets, list):
        if class_type in LOADER_NODES:
            refs.extend(_extract_from_widgets(class_type, widgets))
        else:
            refs.extend(_extract_generic_widgets(class_type, widgets))

    inputs = node.get("inputs")
    if isinstance(inputs, dict) and class_type in LOADER_NODES:
        for field_name, category in LOADER_NODES[class_type]:
            value = inputs.get(field_name)
            if _is_model_filename(value):
                refs.append((str(value), category))

    return refs


def _extract_from_api_format(data: dict[str, Any]) -> list[tuple[str, str]]:
    refs: list[tuple[str, str]] = []
    for _node_id, node in data.items():
        if not isinstance(node, dict):
            continue
        class_type = node.get("class_type")
        if not isinstance(class_type, str):
            continue
        if class_type not in LOADER_NODES:
            continue
        inputs = node.get("inputs")
        if not isinstance(inputs, dict):
            continue
        for field_name, category in LOADER_NODES[class_type]:
            value = inputs.get(field_name)
            if _is_model_filename(value):
                refs.append((str(value), category))
    return refs


def _extract_from_ui_format(data: dict[str, Any]) -> list[tuple[str, str]]:
    refs: list[tuple[str, str]] = []
    seen: set[tuple[str, str]] = set()
    for node in _iter_workflow_nodes(data):
        for ref in _extract_from_node(node):
            if ref not in seen:
                seen.add(ref)
                refs.append(ref)
    return refs


def parse_workflow_data(data: dict[str, Any]) -> list[tuple[str, str]]:
    if "nodes" in data:
        return _extract_from_ui_format(data)
    return _extract_from_api_format(data)


def parse_workflow_file(path: Path, workflows_dir: Path | None = None) -> list[tuple[str, str]]:
    """Return list of (filename, category) referenced in a workflow file."""
    wf_dir = workflows_dir or path.parent
    data = _read_workflow_json(path, wf_dir)
    if not data:
        return []
    return parse_workflow_data(data)


def find_model_file(models_root: Path, filename: str) -> Path | None:
    """Find a model file by basename anywhere under models_root."""
    if not models_root.is_dir():
        return None
    matches: list[Path] = []
    for path in models_root.rglob(filename):
        if path.is_file() and path.name == filename:
            matches.append(path)
    if not matches:
        return None
    return min(matches, key=lambda p: len(p.parts))


def scan_workflows(workflows_dir: Path) -> tuple[list[ModelRef], int]:
    """Scan all JSON workflows and aggregate model references."""
    if not workflows_dir.is_dir():
        return [], 0

    aggregated: dict[tuple[str, str], ModelRef] = {}
    file_count = 0
    resolved_dir = workflows_dir.resolve()
    for path in sorted(resolved_dir.rglob("*.json")):
        if not path.is_file():
            continue
        file_count += 1
        wf_name = path.stem
        for filename, category in parse_workflow_file(path, resolved_dir):
            key = (filename, category)
            if key not in aggregated:
                aggregated[key] = ModelRef(filename=filename, category=category)
            if wf_name not in aggregated[key].workflows:
                aggregated[key].workflows.append(wf_name)

    return sorted(aggregated.values(), key=lambda r: (r.category, r.filename)), file_count