import logging
from dataclasses import dataclass
from typing import Any, Generator, Set

import pandas as pd
from pydantic import BaseModel, Field, PositiveFloat, PositiveInt

from graph.connections.db_connection import DbConnection

logger = logging.getLogger(__name__)


@dataclass
class ProjectedGraph:
    """Simple data class representing a projected graph for algorithm execution."""

    name: str
    node_groups: Set[str]
    edge_groups: Set[str]


class PageRankParams(BaseModel):
    damping_factor: float = Field(..., ge=0.0, le=1.0)
    max_iterations: PositiveInt
    tolerance: PositiveFloat
    normalizeInitial: bool


def compute_page_rank(
    conn: DbConnection, projected_graph: ProjectedGraph, params: PageRankParams, batch_size: int = 10000
) -> Generator[pd.DataFrame, Any, None]:
    from graph.connections.kuzu.db_connection import KuzuConnection
    from graph.connections.neo4j.db_connection import Neo4jConnection

    if isinstance(conn, KuzuConnection):
        from graph.exploration.algorithms.kuzu.page_rank import compute_page_rank as compute_page_rank_kuzu

        return compute_page_rank_kuzu(conn, projected_graph, params, batch_size)
    elif isinstance(conn, Neo4jConnection):
        from graph.exploration.algorithms.neo4j.page_rank import compute_page_rank as compute_page_rank_neo4j

        return compute_page_rank_neo4j(conn, projected_graph, params, batch_size)
    else:
        raise ValueError(f"Unsupported connection type: {type(conn).__name__}")
