from __future__ import annotations

from app.modules.rag.indexing.code.edges.dataflow_slice_builder import DataflowSlice
from app.modules.rag.indexing.code.edges.extractor import PyEdge
from app.modules.rag.indexing.code.entrypoints.execution_trace_builder import ExecutionTrace
from app.modules.rag.indexing.code.roles.models import SemanticRoleRecord
from app.modules.rag.indexing.code.roles.scorer import SemanticRoleScorer
from app.modules.rag.indexing.code.symbols.extractor import PySymbol


class SemanticRoleBuilder:
    def __init__(self, scorer: SemanticRoleScorer | None = None) -> None:
        self._scorer = scorer or SemanticRoleScorer()

    def build(
        self,
        *,
        symbols: list[PySymbol],
        edges: list[PyEdge],
        dataflow_slices: list[DataflowSlice],
        execution_traces: list[ExecutionTrace],
    ) -> list[SemanticRoleRecord]:
        class_symbols = [symbol for symbol in symbols if symbol.kind == "class"]
        result: list[SemanticRoleRecord] = []
        for symbol in class_symbols:
            class_edges = _class_edges(symbol, edges)
            scored = self._scorer.score(
                symbol=symbol,
                class_edges=class_edges,
                dataflow_slices=dataflow_slices,
                execution_traces=execution_traces,
            )
            if scored is None:
                continue
            result.append(scored)
        return result


def _class_edges(symbol: PySymbol, edges: list[PyEdge]) -> list[PyEdge]:
    scope = symbol.qname
    return [edge for edge in edges if edge.src_qname == scope or edge.src_qname.startswith(f"{scope}.")]
