"""VULPO-style lightweight repository context extraction (spec §4.5).

For each target function, extracts (a) callee function bodies (1 hop),
(b) callee signatures only (2 hops), (c) relevant global/type declarations —
budgeted to <=4096 tokens, truncating callee bodies first, then dropping
2-hop signatures. Where repo access is unavailable, context is empty and
`has_context=False` (spec: "never fabricate context").

Token counting here is a simple whitespace-split approximation, not a real
tokenizer — good enough for budgeting (consistently conservative: real BPE
tokenizers produce >= this many tokens for code, so this under-counts and
stays safely under budget, never over).
"""
from __future__ import annotations

from dataclasses import dataclass, field

DEFAULT_TOKEN_BUDGET = 4096


def approx_token_count(text: str) -> int:
    return len(text.split())


@dataclass
class RepoContext:
    callee_bodies: dict[str, str] = field(default_factory=dict)  # name -> source (1-hop)
    callee_signatures: dict[str, str] = field(default_factory=dict)  # name -> signature (2-hop)
    global_decls: list[str] = field(default_factory=list)
    has_context: bool = False
    truncated: bool = False


def build_context(
    callee_bodies: dict[str, str],
    callee_signatures: dict[str, str],
    global_decls: list[str],
    token_budget: int = DEFAULT_TOKEN_BUDGET,
) -> RepoContext:
    """Assembles context under `token_budget`, truncating callee bodies first
    (dropping lowest-priority ones entirely once budget is exhausted — not
    mid-body, since a half-truncated function body is worse than none), then
    dropping 2-hop signatures (spec §4.5's stated truncation order)."""
    if not callee_bodies and not callee_signatures and not global_decls:
        return RepoContext(has_context=False)

    budget = token_budget
    truncated = False

    kept_decls = []
    for decl in global_decls:
        cost = approx_token_count(decl)
        if cost <= budget:
            kept_decls.append(decl)
            budget -= cost
        else:
            truncated = True

    kept_bodies: dict[str, str] = {}
    for name, body in callee_bodies.items():
        cost = approx_token_count(body)
        if cost <= budget:
            kept_bodies[name] = body
            budget -= cost
        else:
            truncated = True  # drop this callee body entirely, not mid-body

    kept_sigs: dict[str, str] = {}
    for name, sig in callee_signatures.items():
        cost = approx_token_count(sig)
        if cost <= budget:
            kept_sigs[name] = sig
            budget -= cost
        else:
            truncated = True

    return RepoContext(
        callee_bodies=kept_bodies,
        callee_signatures=kept_sigs,
        global_decls=kept_decls,
        has_context=True,
        truncated=truncated,
    )


def render_context(ctx: RepoContext) -> str:
    """Flattens a RepoContext into the text block the prompt template embeds."""
    if not ctx.has_context:
        return ""
    parts: list[str] = []
    if ctx.global_decls:
        parts.append("\n".join(ctx.global_decls))
    for name, body in ctx.callee_bodies.items():
        parts.append(f"// callee: {name}\n{body}")
    for name, sig in ctx.callee_signatures.items():
        parts.append(f"// callee signature: {name}\n{sig}")
    return "\n\n".join(parts)