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.symbols.extractor import PySymbol

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


@dataclass(slots=True)
class DataflowSlice:
    slice_id: str
    path: str
    root_symbol_id: str
    path_symbols: list[str]
    path_symbol_ids: list[str]
    path_edge_types: list[str]
    path_length: int
    start_line: int
    end_line: int


class DataflowSliceBuilder:
    def __init__(
        self,
        *,
        max_path_length: int = 6,
        max_slices_per_attribute: int = 10,
        max_instantiation_prefixes: int = 2,
    ) -> None:
        self._max_path_length = max_path_length
        self._max_slices_per_attribute = max_slices_per_attribute
        self._max_instantiation_prefixes = max_instantiation_prefixes

    def build(self, path: str, edges: list[PyEdge], symbols: list[PySymbol]) -> list[DataflowSlice]:
        if not edges:
            return []
        symbol_ids_by_qname = {symbol.qname: symbol.symbol_id for symbol in symbols}
        writers = _group_by_attribute(edges, edge_type="writes_attr")
        readers = _group_by_attribute(edges, edge_type="reads_attr")
        instantiations = [edge for edge in edges if edge.edge_type == "instantiates"]
        slices: list[DataflowSlice] = []
        seen: set[tuple[tuple[str, ...], tuple[str, ...]]] = set()
        for attribute_key, writer_edges in writers.items():
            reader_edges = readers.get(attribute_key, [])
            if not reader_edges:
                continue
            canonical_attr = _canonical_attribute_ref(writer_edges + reader_edges, attribute_key)
            if _should_ignore(canonical_attr):
                continue
            produced = 0
            for writer in writer_edges:
                if produced >= self._max_slices_per_attribute:
                    break
                for reader in reader_edges:
                    if produced >= self._max_slices_per_attribute:
                        break
                    prefixes = _instantiation_prefixes(
                        instantiations=instantiations,
                        writer=writer,
                        attribute_ref=canonical_attr,
                        limit=self._max_instantiation_prefixes,
                    )
                    for instantiation in prefixes:
                        candidate = _build_slice(
                            path=path,
                            writer=writer,
                            reader=reader,
                            canonical_attr=canonical_attr,
                            instantiation=instantiation,
                            symbol_ids_by_qname=symbol_ids_by_qname,
                            max_path_length=self._max_path_length,
                        )
                        if candidate is None:
                            continue
                        dedupe_key = (tuple(candidate.path_symbols), tuple(candidate.path_edge_types))
                        if dedupe_key in seen:
                            continue
                        seen.add(dedupe_key)
                        slices.append(candidate)
                        produced += 1
                        if produced >= self._max_slices_per_attribute:
                            break
        return slices


def _group_by_attribute(edges: list[PyEdge], *, edge_type: str) -> dict[str, list[PyEdge]]:
    grouped: dict[str, list[PyEdge]] = {}
    for edge in edges:
        if edge.edge_type != edge_type:
            continue
        ref = str(edge.dst_ref or "").strip()
        if not ref:
            continue
        key = ref.rsplit(".", 1)[-1].strip().lower()
        grouped.setdefault(key, []).append(edge)
    return grouped


def _canonical_attribute_ref(edges: list[PyEdge], fallback: str) -> str:
    refs = [str(edge.dst_ref or "").strip() for edge in edges if edge.dst_ref]
    for ref in refs:
        owner = ref.rsplit(".", 1)[0] if "." in ref else ""
        if owner.rsplit(".", 1)[-1][:1].isupper():
            return ref
    for ref in refs:
        if "." in ref:
            return ref
    return fallback


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


def _instantiation_prefixes(
    *,
    instantiations: list[PyEdge],
    writer: PyEdge,
    attribute_ref: str,
    limit: int,
) -> list[PyEdge | None]:
    class_hints = _class_hints(writer, attribute_ref)
    if not class_hints:
        return [None]
    matches: list[PyEdge] = []
    for edge in instantiations:
        dst_ref = str(edge.dst_ref or "")
        tail = dst_ref.rsplit(".", 1)[-1]
        if dst_ref in class_hints or tail in class_hints:
            matches.append(edge)
    if not matches:
        return [None]
    matches.sort(key=lambda item: (item.start_line, item.src_qname))
    return matches[:limit]


def _class_hints(writer: PyEdge, attribute_ref: str) -> set[str]:
    values: set[str] = set()
    owner = attribute_ref.rsplit(".", 1)[0] if "." in attribute_ref else ""
    if owner:
        values.add(owner)
        values.add(owner.rsplit(".", 1)[-1])
    writer_owner = writer.src_qname.split(".", 1)[0]
    if writer_owner[:1].isupper():
        values.add(writer_owner)
    return {item for item in values if item[:1].isupper()}


def _build_slice(
    *,
    path: str,
    writer: PyEdge,
    reader: PyEdge,
    canonical_attr: str,
    instantiation: PyEdge | None,
    symbol_ids_by_qname: dict[str, str],
    max_path_length: int,
) -> DataflowSlice | None:
    path_symbols = [writer.src_qname, canonical_attr, reader.src_qname]
    path_edges = ["writes_attr", "reads_attr"]
    symbol_ids = [writer.src_symbol_id, symbol_ids_by_qname.get(canonical_attr, ""), reader.src_symbol_id]
    start_line = min(writer.start_line, reader.start_line)
    end_line = max(writer.end_line, reader.end_line)
    if instantiation is not None:
        class_name = _class_name(instantiation, writer, canonical_attr)
        path_symbols = [instantiation.src_qname, class_name] + path_symbols
        path_edges = ["instantiates"] + path_edges
        symbol_ids = [
            instantiation.src_symbol_id,
            instantiation.dst_symbol_id or symbol_ids_by_qname.get(class_name, ""),
            *symbol_ids,
        ]
        start_line = min(start_line, instantiation.start_line)
        end_line = max(end_line, instantiation.end_line)
    if len(path_symbols) > max_path_length:
        return None
    if any(_should_ignore(item) for item in path_symbols):
        return None
    slice_id = sha256(
        "|".join(
            [
                path,
                writer.edge_id,
                reader.edge_id,
                instantiation.edge_id if instantiation else "none",
                canonical_attr,
            ]
        ).encode("utf-8")
    ).hexdigest()
    return DataflowSlice(
        slice_id=slice_id,
        path=path,
        root_symbol_id=next((item for item in symbol_ids if item), writer.src_symbol_id),
        path_symbols=path_symbols,
        path_symbol_ids=[item for item in symbol_ids if item],
        path_edge_types=path_edges,
        path_length=len(path_symbols),
        start_line=start_line,
        end_line=end_line,
    )


def _class_name(instantiation: PyEdge, writer: PyEdge, attribute_ref: str) -> str:
    if instantiation.dst_ref:
        return instantiation.dst_ref.rsplit(".", 1)[-1]
    owner = attribute_ref.rsplit(".", 1)[0] if "." in attribute_ref else ""
    if owner:
        return owner.rsplit(".", 1)[-1]
    return writer.src_qname.split(".", 1)[0]
