"""
ChunkHive Relationship Graph
============================

Builds structural relationships between code chunks.

Graph Identity Model:
- Each node is identified by chunk.chunk_id (single source of truth)

This enables:
- Dependency-aware ranking
- Graph-based context expansion
- Structural intelligence
"""

from typing import Dict, List, Set, Any
from collections import defaultdict


class RelationshipGraph:
    """
    Structural graph over SemanticChunks.

    Node identity = chunk.chunk_id
    """

    def __init__(self):
        # chunk_id -> chunk
        self.nodes: Dict[str, Any] = {}

        # adjacency: chunk_id -> related chunk_ids
        self.edges: Dict[str, Set[str]] = defaultdict(set)

        # reverse adjacency
        self.reverse_edges: Dict[str, Set[str]] = defaultdict(set)

    # ============================================================
    # Public API
    # ============================================================

    def build(self, chunks: List[Any]) -> None:
        """
        Build graph from chunk list.
        """

        self.nodes.clear()
        self.edges.clear()
        self.reverse_edges.clear()

        # 1️⃣ Register nodes
        for chunk in chunks:
            if not hasattr(chunk, "chunk_id"):
                continue
            self.nodes[chunk.chunk_id] = chunk

        # 2️⃣ Build relationships
        for chunk in chunks:

            if not hasattr(chunk, "chunk_id"):
                continue

            source_id = chunk.chunk_id

            # ------------------------
            # Function call relations
            # ------------------------
            calls = self._extract_function_calls(chunk)

            for called_name in calls:
                target_id = self._find_chunk_by_name(called_name)
                if target_id:
                    self._add_edge(source_id, target_id)

            # ------------------------
            # Class -> method relations
            # ------------------------
            if getattr(chunk, "chunk_type", None) == "class":
                class_name = getattr(chunk.ast, "name", None)

                for other in chunks:
                    if (
                        getattr(other, "chunk_type", None) == "function"
                        and hasattr(other.ast, "parent")
                        and other.ast.parent == class_name
                        and hasattr(other, "chunk_id")
                    ):
                        self._add_edge(source_id, other.chunk_id)

    def get_related(self, chunk_id: str) -> List[str]:
        """
        Return directly connected chunk_ids.
        """
        return list(self.edges.get(chunk_id, []))

    def get_chunk(self, chunk_id: str):
        """
        Return chunk object by ID.
        """
        return self.nodes.get(chunk_id)

    def get_callers(self, chunk_id: str) -> List[str]:
        """
        Return reverse neighbors (who calls this chunk).
        """
        return list(self.reverse_edges.get(chunk_id, []))

    def expand_context(self, chunk_id: str, depth: int = 1) -> List[Any]:
        """
        Expand context up to given depth.
        """

        visited: Set[str] = set()
        to_visit: Set[str] = {chunk_id}
        results: List[Any] = []

        for _ in range(depth):
            next_level: Set[str] = set()

            for nid in to_visit:
                for neighbor in self.edges.get(nid, []):
                    if neighbor not in visited:
                        visited.add(neighbor)
                        next_level.add(neighbor)
                        if neighbor in self.nodes:
                            results.append(self.nodes[neighbor])

            to_visit = next_level

        return results

    # ============================================================
    # Internal Helpers
    # ============================================================

    def _add_edge(self, source: str, target: str) -> None:
        if source == target:
            return

        self.edges[source].add(target)
        self.reverse_edges[target].add(source)

    def _extract_function_calls(self, chunk: Any) -> Set[str]:
        """
        Naive function call extraction.
        Can be upgraded later with AST-level parsing.
        """

        calls = set()

        if not hasattr(chunk, "code") or not chunk.code:
            return calls

        lines = chunk.code.splitlines()

        for line in lines:
            line = line.strip()

            if "(" in line and ")" in line:
                name = line.split("(")[0].strip()

                if (
                    name.isidentifier()
                    and not name.startswith(("if", "for", "while"))
                ):
                    calls.add(name)

        return calls

    def _find_chunk_by_name(self, name: str) -> str | None:
        """
        Resolve chunk_id by AST name.
        """

        for chunk_id, chunk in self.nodes.items():
            if hasattr(chunk.ast, "name") and chunk.ast.name == name:
                return chunk_id

        return None
