services/knowledge_graph_builder.py

"""Repository Knowledge Graph Builder Service.

Composer that orchestrates the construction of the Repository Knowledge Graph
via extensible providers and returns the read-only Pydantic representation.
"""

import logging
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Set
import networkx as nx

from models.knowledge_graph import (
    KnowledgeGraph,
    KnowledgeGraphNode,
    KnowledgeGraphEdge,
    KnowledgeGraphSummary,
)

logger = logging.getLogger(__name__)


class KnowledgeGraphProvider(ABC):
    """Abstract base class representing a plugin provider for knowledge graph entities/edges."""

    @abstractmethod
    def populate(self, repo_name: str, twin: Any, graph: nx.DiGraph) -> None:
        """Populate entities and relationships into the unified NetworkX graph."""
        pass


class MetadataProvider(KnowledgeGraphProvider):
    """Provides Repository root, Health, Compliance, and Architecture report nodes and edges."""

    def populate(self, repo_name: str, twin: Any, graph: nx.DiGraph) -> None:
        # 1. Add Repository Root node
        repo_id = repo_name
        graph.add_node(
            repo_id,
            type="repository",
            name=repo_name.split("/")[-1],
            owner=repo_name.split("/")[0] if "/" in repo_name else "",
            tech_stack=twin.metadata.get("tech_stack", []),
            total_loc=twin.metadata.get("total_loc", 0),
        )

        # 2. Add Health Report node
        health_id = f"{repo_id}::health"
        graph.add_node(
            health_id,
            type="health",
            overall_score=twin.health_summary.get("overall_score", 0.0),
            grade=twin.health_summary.get("grade", "F"),
            breakdown=twin.health_summary.get("breakdown", {}),
        )
        graph.add_edge(repo_id, health_id, type="HAS_HEALTH")

        # 3. Add Compliance Report node
        compliance_id = f"{repo_id}::compliance"
        graph.add_node(
            compliance_id,
            type="compliance",
            status=twin.compliance_summary.get("status", "non-compliant"),
            reasons=twin.compliance_summary.get("reasons", []),
            dead_code_ratio=twin.compliance_summary.get("dead_code_ratio", 0.0),
        )
        graph.add_edge(repo_id, compliance_id, type="HAS_COMPLIANCE")

        # 4. Add Architecture Component nodes
        arch_summary = twin.architecture_summary
        if arch_summary:
            arch_id = f"{repo_id}::architecture"
            graph.add_node(
                arch_id,
                type="architecture",
                cycles_count=arch_summary.get("cycles_count", 0),
                scc_count=arch_summary.get("strongly_connected_components", 0),
            )
            graph.add_edge(repo_id, arch_id, type="HAS_ARCHITECTURE")


class HierarchyProvider(KnowledgeGraphProvider):
    """Provides Directory and File node entities and CONTAINS relationships."""

    def populate(self, repo_name: str, twin: Any, graph: nx.DiGraph) -> None:
        repo_id = repo_name
        files = twin.files

        added_dirs: Set[str] = set()

        for file_path in files:
            # Normalize path separators
            norm_path = file_path.replace("\\", "/")
            parts = norm_path.split("/")

            # Add directory nodes recursively
            current_dir_id = repo_id
            for i in range(len(parts) - 1):
                sub_dir_path = "/".join(parts[: i + 1])
                sub_dir_id = f"{repo_id}::{sub_dir_path}"

                if sub_dir_id not in added_dirs:
                    graph.add_node(
                        sub_dir_id,
                        type="directory",
                        name=parts[i],
                        path=sub_dir_path,
                    )
                    # Link parent directory/repo to this directory
                    graph.add_edge(current_dir_id, sub_dir_id, type="CONTAINS")
                    added_dirs.add(sub_dir_id)

                current_dir_id = sub_dir_id

            # Add File node
            file_id = f"{repo_id}::{norm_path}"
            graph.add_node(
                file_id,
                type="file",
                name=parts[-1],
                path=norm_path,
            )
            # Link final directory containing it to the file node
            graph.add_edge(current_dir_id, file_id, type="CONTAINS")


class SymbolProvider(KnowledgeGraphProvider):
    """Provides Symbol nodes and DECLARES relationships."""

    def __init__(self, symbol_service: Optional[Any] = None) -> None:
        self.symbol_service = symbol_service

    def populate(self, repo_name: str, twin: Any, graph: nx.DiGraph) -> None:
        repo_id = repo_name
        ss = self.symbol_service
        if ss is None:
            from services.symbol_service import SymbolService

            ss = SymbolService()

        symbol_index = ss.load(repo_name)
        if not symbol_index:
            return

        # Add all symbols
        for sym in symbol_index.symbols:
            norm_file = sym.file_path.replace("\\", "/")
            file_id = f"{repo_id}::{norm_file}"
            qualified_name = (
                f"{sym.parent_class}.{sym.name}" if sym.parent_class else sym.name
            )
            symbol_id = f"{repo_id}::{norm_file}::{qualified_name}"

            graph.add_node(
                symbol_id,
                type="symbol",
                name=sym.name,
                symbol_type=sym.type,
                line_number=sym.line_number,
                language=sym.language,
            )

            # Link File to Symbol declaration
            graph.add_edge(file_id, symbol_id, type="DECLARES")

            # Link Class to Method if nested
            if sym.parent_class:
                class_id = f"{repo_id}::{norm_file}::{sym.parent_class}"
                if class_id in graph:
                    graph.add_edge(class_id, symbol_id, type="DECLARES")


class DependencyProvider(KnowledgeGraphProvider):
    """Provides IMPORTS edges between File nodes based on imports dependency graph."""

    def __init__(self, graph_service: Optional[Any] = None) -> None:
        self.graph_service = graph_service

    def populate(self, repo_name: str, twin: Any, graph: nx.DiGraph) -> None:
        repo_id = repo_name
        gs = self.graph_service
        if gs is None:
            from services.graph_service import GraphService

            gs = GraphService()

        dep_graph = gs.load_graph(repo_name)
        if dep_graph is None:
            return

        for u, v in dep_graph.edges():
            u_norm = u.replace("\\", "/")
            v_norm = v.replace("\\", "/")
            u_id = f"{repo_id}::{u_norm}"
            v_id = f"{repo_id}::{v_norm}"

            # Only add edge if both nodes are present in the hierarchy
            if u_id in graph and v_id in graph:
                graph.add_edge(u_id, v_id, type="IMPORTS")


class CallGraphProvider(KnowledgeGraphProvider):
    """Provides CALLS edges between Symbol nodes based on call graph edges."""

    def __init__(self, graph_service: Optional[Any] = None) -> None:
        self.graph_service = graph_service

    def populate(self, repo_name: str, twin: Any, graph: nx.DiGraph) -> None:
        repo_id = repo_name
        gs = self.graph_service
        if gs is None:
            from services.graph_service import GraphService

            gs = GraphService()

        call_graph = gs.load_graph(f"{repo_name}_call_graph")
        if call_graph is None:
            return

        for u, v in call_graph.edges():
            # Call graph node format: "file_path::qualified_name"
            if "::" in u and "::" in v:
                u_file, u_qual = u.split("::", 1)
                v_file, v_qual = v.split("::", 1)

                u_file_clean = u_file.replace("\\", "/")
                v_file_clean = v_file.replace("\\", "/")
                u_id = f"{repo_id}::{u_file_clean}::{u_qual}"
                v_id = f"{repo_id}::{v_file_clean}::{v_qual}"

                if u_id in graph and v_id in graph:
                    graph.add_edge(u_id, v_id, type="CALLS")


class RepositoryKnowledgeGraphBuilder:
    """Thin composer that runs registered providers to build the RepositoryKnowledgeGraph."""

    def __init__(
        self,
        twin_builder: Optional[Any] = None,
        cache: Optional[Any] = None,
        providers: Optional[List[KnowledgeGraphProvider]] = None,
        symbol_service: Optional[Any] = None,
        graph_service: Optional[Any] = None,
    ) -> None:
        """Initialise the Knowledge Graph composer service."""
        import sys

        backend_deps = sys.modules.get("backend.dependencies")
        if backend_deps:
            symbol_service = symbol_service or getattr(
                backend_deps, "symbol_service", None
            )
            graph_service = graph_service or getattr(
                backend_deps, "graph_service", None
            )
            twin_builder = twin_builder or getattr(
                backend_deps, "repository_twin_builder", None
            )
            cache = cache or getattr(backend_deps, "analysis_cache", None)

        self.twin_builder = twin_builder
        self.cache = cache

        # Set default providers
        self.providers = (
            providers
            if providers is not None
            else [
                MetadataProvider(),
                HierarchyProvider(),
                SymbolProvider(symbol_service),
                DependencyProvider(graph_service),
                CallGraphProvider(graph_service),
            ]
        )

    def register_provider(self, provider: KnowledgeGraphProvider) -> None:
        """Dynamically registers a new plugin provider to extend the Knowledge Graph."""
        self.providers.append(provider)

    def build_networkx_graph(self, repo_name: str) -> nx.DiGraph:
        """Builds and returns the internal NetworkX DiGraph by running all providers."""
        twin = self.twin_builder.build_twin(repo_name)
        graph = nx.DiGraph()

        for provider in self.providers:
            try:
                provider.populate(repo_name, twin, graph)
            except Exception as e:
                logger.error(
                    "Provider %s failed during twin composition for %s: %s",
                    provider.__class__.__name__,
                    repo_name,
                    e,
                    exc_info=True,
                )

        return graph

    def build_graph(self, repo_name: str) -> KnowledgeGraph:
        """Assembles and returns the fully serialized Pydantic KnowledgeGraph."""
        # 1. Try to fetch from cache first
        cached = self.cache.get(repo_name, "knowledge_graph", 1)
        if cached is not None:
            return cached

        # 2. Build internal graph
        nx_graph = self.build_networkx_graph(repo_name)

        # 3. Serialize nodes and edges
        nodes = []
        for node_id, data in nx_graph.nodes(data=True):
            node_type = data.get("type", "unknown")
            props = {k: v for k, v in data.items() if k != "type"}
            nodes.append(
                KnowledgeGraphNode(id=node_id, type=node_type, properties=props)
            )

        edges = []
        for u, v, data in nx_graph.edges(data=True):
            edge_type = data.get("type", "unknown")
            props = {k: v for k, v in data.items() if k != "type"}
            edges.append(
                KnowledgeGraphEdge(source=u, target=v, type=edge_type, properties=props)
            )

        kg = KnowledgeGraph(repository_name=repo_name, nodes=nodes, edges=edges)

        # 4. Save to cache
        self.cache.set(repo_name, "knowledge_graph", kg, 1)

        return kg

    def build_graph_summary(self, repo_name: str) -> KnowledgeGraphSummary:
        """Retrieves or builds a lightweight summary stats model of the Knowledge Graph."""
        kg = self.build_graph(repo_name)

        # Aggregate counts
        node_breakdown: Dict[str, int] = {}
        for node in kg.nodes:
            node_breakdown[node.type] = node_breakdown.get(node.type, 0) + 1

        edge_breakdown: Dict[str, int] = {}
        for edge in kg.edges:
            edge_breakdown[edge.type] = edge_breakdown.get(edge.type, 0) + 1

        return KnowledgeGraphSummary(
            repository_name=repo_name,
            nodes_count=len(kg.nodes),
            edges_count=len(kg.edges),
            node_types_breakdown=node_breakdown,
            edge_types_breakdown=edge_breakdown,
        )