from __future__ import annotations

from collections.abc import Iterable

from .models import Classification, ProcessGroup, ProcessNode, ProcessSample


def build_nodes(
    samples: list[ProcessSample],
    classifications: dict[str, Classification],
) -> dict[str, ProcessNode]:
    nodes = {
        sample.process_key: ProcessNode(
            sample=sample,
            classification=classifications.get(sample.process_key, Classification("unknown", 0.0)),
        )
        for sample in samples
    }
    pid_to_key = {sample.pid: sample.process_key for sample in samples}

    for node in nodes.values():
        parent_key = pid_to_key.get(node.sample.ppid)
        if parent_key and parent_key in nodes and parent_key != node.process_key:
            node.parent_key = parent_key
            nodes[parent_key].children.append(node)

    for node in nodes.values():
        node.children.sort(key=lambda item: item.sample.pid)

    for root in _roots(nodes.values()):
        _propagate_agent_context(root, inherited=False)

    return nodes


def _roots(nodes: Iterable[ProcessNode]) -> list[ProcessNode]:
    node_list = list(nodes)
    keys = {node.process_key for node in node_list}
    return [node for node in node_list if not node.parent_key or node.parent_key not in keys]


def _propagate_agent_context(node: ProcessNode, *, inherited: bool) -> None:
    current = inherited or node.classification.is_agent or node.classification.is_agent_related
    node.inherited_agent = current
    for child in node.children:
        _propagate_agent_context(child, inherited=current)


def iter_descendants(node: ProcessNode, *, include_self: bool = True) -> Iterable[ProcessNode]:
    if include_self:
        yield node
    for child in node.children:
        yield from iter_descendants(child, include_self=True)


def _has_agent_ancestor(node: ProcessNode, nodes: dict[str, ProcessNode]) -> bool:
    parent_key = node.parent_key
    while parent_key and parent_key in nodes:
        parent = nodes[parent_key]
        if parent.classification.is_agent:
            return True
        parent_key = parent.parent_key
    return False


def build_groups(nodes: dict[str, ProcessNode]) -> list[ProcessGroup]:
    group_roots: list[ProcessNode] = []
    for node in nodes.values():
        classification = node.classification
        if classification.is_agent and not _has_agent_ancestor(node, nodes):
            group_roots.append(node)
        elif (
            classification.is_mcp or classification.is_bridge
        ) and not _has_agent_ancestor(node, nodes):
            group_roots.append(node)

    seen: set[str] = set()
    groups: list[ProcessGroup] = []
    for root in sorted(group_roots, key=lambda item: item.sample.pid):
        if root.process_key in seen:
            continue
        seen.add(root.process_key)
        members = list(iter_descendants(root))
        bridge_count = sum(1 for node in members if node.classification.is_bridge)
        mcp_count = sum(1 for node in members if node.classification.is_mcp)
        groups.append(
            ProcessGroup(
                group_id=root.process_key,
                root=root,
                nodes=members,
                label=_group_label(root),
                child_count=max(0, len(members) - 1),
                bridge_count=bridge_count,
                mcp_count=mcp_count,
                total_rss_bytes=sum(node.sample.rss_bytes for node in members),
                max_age_seconds=max((node.sample.age_seconds for node in members), default=0.0),
            )
        )
    return groups


def _group_label(root: ProcessNode) -> str:
    argv0 = root.sample.argv[0] if root.sample.argv else root.sample.comm
    name = argv0.rsplit("/", 1)[-1] or root.sample.comm
    return f"{name} - PID {root.sample.pid}"