DocForge2/src/docforge/context.py

"""Deterministic bounded context compilation with explicit selection provenance."""

from __future__ import annotations

from collections import deque

from .errors import DocForgeError
from .index import ProjectIndex
from .models import ContextEntry, ContextProfile, Node, ProjectSnapshot


def _estimate_tokens(text: str) -> int:
    return max(1, (len(text) + 3) // 4)


def _node_text(node: Node) -> str:
    return (
        f"ID: {node.node_id}\nTitle: {node.title}\nFamily: {node.family}\n"
        f"Authority: {node.authority}\nStatus: {node.status}\nSource: {node.source_path}\n"
        f"Summary: {node.summary}\n\n{node.content}"
    )


def _profile(snapshot: ProjectSnapshot, profile_id: str) -> ContextProfile:
    matches = [
        profile for profile in snapshot.descriptor.profiles if profile.profile_id == profile_id
    ]
    if len(matches) != 1:
        raise DocForgeError(
            "missing_profile", "No context profile has the requested ID", id=profile_id
        )
    return matches[0]


def compile_context(
    index: ProjectIndex, profile_id: str, budget: int | None = None
) -> dict[str, object]:
    def select(snapshot: ProjectSnapshot) -> dict[str, object]:
        profile = _profile(snapshot, profile_id)
        selected_budget = profile.token_budget if budget is None else budget
        if (
            isinstance(selected_budget, bool)
            or selected_budget < 1
            or selected_budget > snapshot.descriptor.limits.max_context_tokens
        ):
            raise DocForgeError("invalid_budget", "Context budget is outside the configured range")

        node_by_id = {node.node_id: node for node in snapshot.nodes}
        dependency_lists: dict[str, list[str]] = {node_id: [] for node_id in node_by_id}
        for edge in snapshot.edges:
            if edge.relation == "depends_on":
                dependency_lists[edge.source_id].append(edge.target_id)
        dependency_edges = {
            node_id: tuple(targets) for node_id, targets in dependency_lists.items()
        }
        reasons: dict[str, str] = {
            node_id: "required by profile" for node_id in profile.required_nodes
        }
        queue = deque((node_id, 0) for node_id in profile.required_nodes)
        while queue:
            node_id, depth = queue.popleft()
            if depth >= profile.dependency_depth:
                continue
            for dependency in dependency_edges[node_id]:
                if dependency not in reasons:
                    reasons[dependency] = f"dependency of {node_id}"
                    queue.append((dependency, depth + 1))

        eligible = [
            node
            for node in snapshot.nodes
            if (not profile.families or node.family in profile.families)
            and (not profile.statuses or node.status in profile.statuses)
        ]
        ordered_ids = [*profile.required_nodes]
        ordered_ids.extend(sorted(set(reasons) - set(ordered_ids)))
        ordered_ids.extend(node.node_id for node in eligible if node.node_id not in reasons)

        entries: list[ContextEntry] = []
        omissions: list[dict[str, str]] = []
        used_tokens = 0
        required = set(profile.required_nodes)
        for node_id in ordered_ids:
            node = node_by_id[node_id]
            text = _node_text(node)
            tokens = _estimate_tokens(text)
            if used_tokens + tokens > selected_budget:
                if node_id in required:
                    raise DocForgeError(
                        "budget_too_small",
                        "Context budget cannot contain every required node",
                        node_id=node_id,
                        required_tokens=used_tokens + tokens,
                    )
                omissions.append({"node_id": node_id, "reason": "token budget"})
                continue
            entries.append(
                ContextEntry(
                    node_id=node_id,
                    reason=reasons.get(node_id, "eligible profile node"),
                    estimated_tokens=tokens,
                    source_path=node.source_path,
                    content_hash=node.content_hash,
                    text=text,
                )
            )
            used_tokens += tokens

        return {
            "profile": profile.profile_id,
            "budget": selected_budget,
            "estimated_tokens": used_tokens,
            "entries": [entry.as_dict() for entry in entries],
            "omissions": omissions,
        }

    return index.read_project_snapshot(select)