"""Model registry loading: read each workflow, validate links/tokens, and
decide which model an endpoint may run.

Startup validation of a workflow is local here (broken links + token check);
the collected graphs are then dry-injected and sent to the server elsewhere
(see main.lifespan).
"""
import json
from dataclasses import dataclass, field
from pathlib import Path

from .config import Settings
from .errors import AdapterError, ConfigError, InvalidRequest
from .injector import extract_tokens


class TemplateError(AdapterError):
    """Malformed workflow template; blocks startup."""

    def __init__(self, message: str):
        super().__init__(500, message, "server_error", "template_error")


@dataclass(frozen=True)
class Model:
    id: str
    graph: dict
    declared: frozenset[str]
    owned_by: str = "comfy-adapter"
    workflow: Path = field(default=Path("."))


# Token sets belong to the endpoint, not the model. A model is valid for an
# endpoint iff required(endpoint) <= declared(graph) <= supplied(endpoint).
ENDPOINT_VOCAB: dict[str, dict[str, frozenset[str]]] = {
    "generations": {
        "supplied": frozenset({"prompt", "negative", "seed", "width", "height"}),
        "required": frozenset({"prompt"}),
    },
    "edits": {
        "supplied": frozenset({"prompt", "negative", "seed", "image"}),
        "required": frozenset({"prompt", "image"}),
    },
}


def find_broken_links(graph: dict) -> list[str]:
    """References to node ids that do not exist in the graph."""
    errors = []
    for node_id, node in graph.items():
        for field, value in node.get("inputs", {}).items():
            if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str):
                if value[0] not in graph:
                    errors.append(f"node {node_id}: input '{field}' links to missing node '{value[0]}'")
    return errors


def is_compatible(declared: frozenset[str], endpoint: str) -> bool:
    vocab = ENDPOINT_VOCAB[endpoint]
    return vocab["required"] <= declared and declared <= vocab["supplied"]


def select_model(models: dict[str, Model], model: str | None, endpoint: str) -> Model:
    """Pick the model for a request by id, or the first compatible one.

    An explicit id that is unknown or incompatible with the endpoint is a 400,
    as is a missing id when no registered model is compatible.
    """
    if model is not None:
        found = models.get(model)
        if found is None:
            raise InvalidRequest(f"unknown model '{model}'")
        if not is_compatible(found.declared, endpoint):
            raise InvalidRequest(f"model '{model}' is not compatible with this endpoint")
        return found
    for candidate in models.values():  # declaration order = priority
        if is_compatible(candidate.declared, endpoint):
            return candidate
    raise InvalidRequest(f"no model in the registry is compatible with this endpoint")


def _read_graph(path: Path, model_id: str) -> dict:
    try:
        with open(path, encoding="utf-8") as f:
            return json.load(f)
    except FileNotFoundError as e:
        raise TemplateError(f"model '{model_id}': workflow file not found: {path}") from e
    except json.JSONDecodeError as e:
        raise TemplateError(f"model '{model_id}': workflow is not valid JSON ({path}): {e}") from e


def load_model(model_id: str, workflow_path: Path, owned_by: str) -> Model:
    graph = _read_graph(workflow_path, model_id)
    errors = find_broken_links(graph)
    declared, unknown = extract_tokens(graph)
    errors += unknown
    if errors:
        raise TemplateError(
            f"workflow for model '{model_id}' ({workflow_path}) is invalid:\n  " + "\n  ".join(errors)
        )
    return Model(id=model_id, graph=graph, declared=frozenset(declared), owned_by=owned_by, workflow=workflow_path)


def load_registry(settings: Settings) -> dict[str, Model]:
    """Load every model in the registry, preserving order for default selection."""
    base_dir = settings.models_file.parent
    registry = settings.load_registry()
    models: dict[str, Model] = {}
    for model_id, entry in registry.items():
        path = entry.workflow if entry.workflow.is_absolute() else base_dir / entry.workflow
        models[model_id] = load_model(model_id, path, entry.owned_by)
    return models