from __future__ import annotations

from dataclasses import dataclass
from hashlib import sha256

from app.modules.rag.indexing.code.edges.extractor import PyEdge
from app.modules.rag.indexing.code.entrypoints.registry import Entrypoint
from app.modules.rag.indexing.code.symbols.extractor import PySymbol


@dataclass(slots=True)
class ExecutionTrace:
    trace_id: str
    path: str
    entry_id: str
    entry_symbol_id: str
    path_symbols: list[str]
    path_symbol_ids: list[str]
    path_length: int
    start_line: int
    end_line: int


class ExecutionTraceBuilder:
    _IGNORE_PREFIXES = ("logger", "logging", "print", "stdlib", "typing", "dataclasses")

    def __init__(self, *, max_trace_depth: int = 8, max_traces_per_entrypoint: int = 20) -> None:
        self._max_trace_depth = max_trace_depth
        self._max_traces_per_entrypoint = max_traces_per_entrypoint

    def build(
        self,
        path: str,
        entrypoints: list[Entrypoint],
        edges: list[PyEdge],
        symbols: list[PySymbol],
    ) -> list[ExecutionTrace]:
        if not entrypoints:
            return []
        symbol_map = {symbol.symbol_id: symbol for symbol in symbols}
        call_map = self._call_edges(edges)
        traces: list[ExecutionTrace] = []
        seen: set[tuple[str, tuple[str, ...]]] = set()
        for entrypoint in entrypoints:
            raw_paths = self._dfs_paths(entrypoint.handler_symbol_id, call_map)
            for symbol_ids in raw_paths:
                if len(symbol_ids) < 2:
                    continue
                key = (entrypoint.entry_id, tuple(symbol_ids))
                if key in seen:
                    continue
                seen.add(key)
                trace = self._trace_for_path(path, entrypoint, symbol_ids, symbol_map)
                if trace is None:
                    continue
                traces.append(trace)
        return traces

    def _call_edges(self, edges: list[PyEdge]) -> dict[str, list[PyEdge]]:
        grouped: dict[str, list[PyEdge]] = {}
        for edge in edges:
            if edge.edge_type != "calls" or not edge.dst_symbol_id:
                continue
            grouped.setdefault(edge.src_symbol_id, []).append(edge)
        for values in grouped.values():
            values.sort(key=lambda item: (item.start_line, item.dst_symbol_id or ""))
        return grouped

    def _dfs_paths(self, start_symbol_id: str, call_map: dict[str, list[PyEdge]]) -> list[list[str]]:
        if not start_symbol_id:
            return []
        results: list[list[str]] = []
        stack: list[tuple[str, list[str], set[str]]] = [(start_symbol_id, [start_symbol_id], {start_symbol_id})]
        while stack and len(results) < self._max_traces_per_entrypoint:
            current, path_ids, visited = stack.pop()
            edges = call_map.get(current, [])
            if not edges or len(path_ids) >= self._max_trace_depth:
                results.append(path_ids)
                continue
            pushed = False
            for edge in reversed(edges):
                dst_symbol_id = edge.dst_symbol_id
                if not dst_symbol_id or dst_symbol_id in visited:
                    continue
                pushed = True
                stack.append((dst_symbol_id, path_ids + [dst_symbol_id], visited | {dst_symbol_id}))
            if not pushed:
                results.append(path_ids)
        return results

    def _trace_for_path(
        self,
        path: str,
        entrypoint: Entrypoint,
        symbol_ids: list[str],
        symbol_map: dict[str, PySymbol],
    ) -> ExecutionTrace | None:
        path_symbols = [self._symbol_name(symbol_map, symbol_id) for symbol_id in symbol_ids]
        if any(self._should_ignore(item) for item in path_symbols):
            return None
        lines = [symbol_map.get(symbol_id) for symbol_id in symbol_ids]
        known_lines = [item for item in lines if item is not None]
        start_line = min([entrypoint.start_line, *[item.start_line for item in known_lines]])
        end_line = max([entrypoint.end_line, *[item.end_line for item in known_lines]])
        trace_id = sha256(f"{path}|{entrypoint.entry_id}|{'|'.join(symbol_ids)}".encode("utf-8")).hexdigest()
        return ExecutionTrace(
            trace_id=trace_id,
            path=path,
            entry_id=entrypoint.entry_id,
            entry_symbol_id=entrypoint.handler_symbol_id,
            path_symbols=path_symbols,
            path_symbol_ids=symbol_ids,
            path_length=len(path_symbols),
            start_line=start_line,
            end_line=end_line,
        )

    def _symbol_name(self, symbol_map: dict[str, PySymbol], symbol_id: str) -> str:
        symbol = symbol_map.get(symbol_id)
        if symbol is None:
            return symbol_id
        return symbol.qname

    def _should_ignore(self, value: str) -> bool:
        lowered = value.lower()
        return any(lowered == item or lowered.startswith(f"{item}.") for item in self._IGNORE_PREFIXES)
