from __future__ import annotations

from typing import Literal

from pydantic import BaseModel, model_validator

from graph.exploration.algorithms.graph_algorithm_registry import (
    GRAPH_FEATURE_ALGORITHM_DEFINITIONS,
    GRAPH_FEATURE_ALGORITHM_NAMES,
)


class GraphFeaturesParams(BaseModel):
    directed_graph: bool
    output_type: Literal["output_nodes", "output_edges"]
    degree: bool = True
    eigenvector_centrality: bool = False
    clustering: bool = False
    triangles: bool = False
    closeness: bool = False
    pagerank: bool = False
    sq_clustering: bool = False
    connected_components: bool = False

    @model_validator(mode="after")
    def validate_directed_graph_feature_support(self) -> "GraphFeaturesParams":
        if not selected_graph_feature_names(self):
            raise ValueError("At least one graph feature must be selected.")

        unsupported_features = (
            tuple(
                definition.name
                for definition in GRAPH_FEATURE_ALGORITHM_DEFINITIONS
                if not definition.supports_directed and getattr(self, definition.name)
            )
            if self.directed_graph
            else ()
        )
        if unsupported_features:
            feature_names = ", ".join(unsupported_features)
            raise ValueError(f"Directed graph feature requests do not support: {feature_names}.")
        return self


def selected_graph_feature_names(params: GraphFeaturesParams) -> tuple[str, ...]:
    return tuple(algorithm_name for algorithm_name in GRAPH_FEATURE_ALGORITHM_NAMES if getattr(params, algorithm_name))


__all__ = ["GraphFeaturesParams", "selected_graph_feature_names"]
