ATOCore/src/atocore/context/builder.py

"""Context pack assembly: retrieve, rank, budget, format.

Trust precedence (per Master Plan):
  1. Trusted Project State → always included first, highest authority
  2. Identity + Preference memories → included next
  3. Retrieved chunks → ranked, deduplicated, budget-constrained
"""

import time
from dataclasses import dataclass, field
from pathlib import Path

import atocore.config as _config
from atocore.context.project_state import format_project_state, get_state
from atocore.memory.service import get_memories_for_context
from atocore.observability.logger import get_logger
from atocore.projects.registry import get_registered_project
from atocore.retrieval.retriever import ChunkResult, retrieve

log = get_logger("context_builder")

SYSTEM_PREFIX = (
    "You have access to the following personal context from the user's knowledge base.\n"
    "Use it to inform your answer. If the context is not relevant, ignore it.\n"
    "Do not mention the context system unless asked.\n"
    "When project state is provided, treat it as the most authoritative source."
)

# Budget allocation (per Master Plan section 9):
#   identity: 5%, preferences: 5%, project state: 20%, retrieval: 60%+
PROJECT_STATE_BUDGET_RATIO = 0.20
MEMORY_BUDGET_RATIO = 0.10  # 5% identity + 5% preference

# Last built context pack for debug inspection
_last_context_pack: "ContextPack | None" = None


@dataclass
class ContextChunk:
    content: str
    source_file: str
    heading_path: str
    score: float
    char_count: int


@dataclass
class ContextPack:
    chunks_used: list[ContextChunk] = field(default_factory=list)
    project_state_text: str = ""
    project_state_chars: int = 0
    memory_text: str = ""
    memory_chars: int = 0
    total_chars: int = 0
    budget: int = 0
    budget_remaining: int = 0
    formatted_context: str = ""
    full_prompt: str = ""
    query: str = ""
    project_hint: str = ""
    duration_ms: int = 0


def build_context(
    user_prompt: str,
    project_hint: str | None = None,
    budget: int | None = None,
) -> ContextPack:
    """Build a context pack for a user prompt.

    Trust precedence applied:
      1. Project state is injected first (highest trust)
      2. Identity + preference memories (second trust level)
      3. Retrieved chunks fill the remaining budget
    """
    global _last_context_pack
    start = time.time()
    budget = _config.settings.context_budget if budget is None else max(budget, 0)

    # 1. Get Trusted Project State (highest precedence)
    project_state_text = ""
    project_state_chars = 0
    project_state_budget = min(
        budget,
        max(0, int(budget * PROJECT_STATE_BUDGET_RATIO)),
    )

    # Resolve the project hint through the registry so callers can pass
    # an alias (`p05`, `gigabit`) and still find trusted state stored
    # under the canonical project id (`p05-interferometer`,
    # `p04-gigabit`). The retriever already does this for the
    # project-match boost — the project_state lookup needs the same
    # courtesy. If the registry has no entry for the hint, fall back to
    # the raw hint so a hand-curated project_state entry that predates
    # the registry still works.
    canonical_project = project_hint
    if project_hint:
        registered = get_registered_project(project_hint)
        if registered is not None:
            canonical_project = registered.project_id

        state_entries = get_state(canonical_project)
        if state_entries:
            project_state_text = format_project_state(state_entries)
            project_state_text, project_state_chars = _truncate_text_block(
                project_state_text,
                project_state_budget or budget,
            )

    # 2. Get identity + preference memories (second precedence)
    memory_budget = min(int(budget * MEMORY_BUDGET_RATIO), max(budget - project_state_chars, 0))
    memory_text, memory_chars = get_memories_for_context(
        memory_types=["identity", "preference"],
        budget=memory_budget,
    )

    # 3. Calculate remaining budget for retrieval
    retrieval_budget = budget - project_state_chars - memory_chars

    # 4. Retrieve candidates
    candidates = (
        retrieve(
            user_prompt,
            top_k=_config.settings.context_top_k,
            project_hint=project_hint,
        )
        if retrieval_budget > 0
        else []
    )

    # 5. Score and rank
    scored = _rank_chunks(candidates, project_hint)

    # 6. Select within remaining budget
    selected = _select_within_budget(scored, max(retrieval_budget, 0))

    # 7. Format full context
    formatted = _format_full_context(project_state_text, memory_text, selected)
    if len(formatted) > budget:
        formatted, selected = _trim_context_to_budget(
            project_state_text,
            memory_text,
            selected,
            budget,
        )

    # 8. Build full prompt
    full_prompt = f"{SYSTEM_PREFIX}\n\n{formatted}\n\n{user_prompt}"

    project_state_chars = len(project_state_text)
    memory_chars = len(memory_text)
    retrieval_chars = sum(c.char_count for c in selected)
    total_chars = len(formatted)
    duration_ms = int((time.time() - start) * 1000)

    pack = ContextPack(
        chunks_used=selected,
        project_state_text=project_state_text,
        project_state_chars=project_state_chars,
        memory_text=memory_text,
        memory_chars=memory_chars,
        total_chars=total_chars,
        budget=budget,
        budget_remaining=budget - total_chars,
        formatted_context=formatted,
        full_prompt=full_prompt,
        query=user_prompt,
        project_hint=project_hint or "",
        duration_ms=duration_ms,
    )

    _last_context_pack = pack

    log.info(
        "context_built",
        chunks_used=len(selected),
        project_state_chars=project_state_chars,
        memory_chars=memory_chars,
        retrieval_chars=retrieval_chars,
        total_chars=total_chars,
        budget_remaining=budget - total_chars,
        duration_ms=duration_ms,
    )
    log.debug("context_pack_detail", pack=_pack_to_dict(pack))

    return pack


def get_last_context_pack() -> ContextPack | None:
    """Return the last built context pack for debug inspection."""
    return _last_context_pack


def _rank_chunks(
    candidates: list[ChunkResult],
    project_hint: str | None,
) -> list[tuple[float, ChunkResult]]:
    """Rank candidates with boosting for project match."""
    scored = []
    seen_content: set[str] = set()

    for chunk in candidates:
        # Deduplicate by content prefix (first 200 chars)
        content_key = chunk.content[:200]
        if content_key in seen_content:
            continue
        seen_content.add(content_key)

        # Base score from similarity
        final_score = chunk.score

        # Project boost
        if project_hint:
            tags_str = chunk.tags.lower() if chunk.tags else ""
            source_str = chunk.source_file.lower()
            title_str = chunk.title.lower() if chunk.title else ""
            hint_lower = project_hint.lower()

            if hint_lower in tags_str or hint_lower in source_str or hint_lower in title_str:
                final_score *= 1.3

        scored.append((final_score, chunk))

    # Sort by score descending
    scored.sort(key=lambda x: x[0], reverse=True)
    return scored


def _select_within_budget(
    scored: list[tuple[float, ChunkResult]],
    budget: int,
) -> list[ContextChunk]:
    """Select top chunks that fit within the character budget."""
    selected = []
    used = 0

    for score, chunk in scored:
        chunk_len = len(chunk.content)
        if used + chunk_len > budget:
            continue
        selected.append(
            ContextChunk(
                content=chunk.content,
                source_file=_shorten_path(chunk.source_file),
                heading_path=chunk.heading_path,
                score=score,
                char_count=chunk_len,
            )
        )
        used += chunk_len

    return selected


def _format_full_context(
    project_state_text: str,
    memory_text: str,
    chunks: list[ContextChunk],
) -> str:
    """Format project state + memories + retrieved chunks into full context block."""
    parts = []

    # 1. Project state first (highest trust)
    if project_state_text:
        parts.append(project_state_text)
        parts.append("")

    # 2. Identity + preference memories (second trust level)
    if memory_text:
        parts.append(memory_text)
        parts.append("")

    # 3. Retrieved chunks (lowest trust)
    if chunks:
        parts.append("--- AtoCore Retrieved Context ---")
        if project_state_text:
            parts.append("If retrieved context conflicts with Trusted Project State above, trust the Trusted Project State.")
        for chunk in chunks:
            parts.append(
                f"[Source: {chunk.source_file} | Section: {chunk.heading_path} | Score: {chunk.score:.2f}]"
            )
            parts.append(chunk.content)
            parts.append("")
        parts.append("--- End Context ---")
    elif not project_state_text and not memory_text:
        parts.append("--- AtoCore Context ---\nNo relevant context found.\n--- End Context ---")

    return "\n".join(parts)


def _shorten_path(path: str) -> str:
    """Shorten an absolute path to a relative-like display."""
    p = Path(path)
    parts = p.parts
    if len(parts) > 3:
        return str(Path(*parts[-3:]))
    return str(p)


def _pack_to_dict(pack: ContextPack) -> dict:
    """Convert a context pack to a JSON-serializable dict."""
    return {
        "query": pack.query,
        "project_hint": pack.project_hint,
        "project_state_chars": pack.project_state_chars,
        "memory_chars": pack.memory_chars,
        "chunks_used": len(pack.chunks_used),
        "total_chars": pack.total_chars,
        "budget": pack.budget,
        "budget_remaining": pack.budget_remaining,
        "duration_ms": pack.duration_ms,
        "has_project_state": bool(pack.project_state_text),
        "has_memories": bool(pack.memory_text),
        "chunks": [
            {
                "source_file": c.source_file,
                "heading_path": c.heading_path,
                "score": c.score,
                "char_count": c.char_count,
                "content_preview": c.content[:100],
            }
            for c in pack.chunks_used
        ],
    }


def _truncate_text_block(text: str, budget: int) -> tuple[str, int]:
    """Trim a formatted text block so trusted tiers cannot exceed the total budget."""
    if budget <= 0 or not text:
        return "", 0
    if len(text) <= budget:
        return text, len(text)
    if budget <= 3:
        trimmed = text[:budget]
    else:
        trimmed = f"{text[: budget - 3].rstrip()}..."
    return trimmed, len(trimmed)


def _trim_context_to_budget(
    project_state_text: str,
    memory_text: str,
    chunks: list[ContextChunk],
    budget: int,
) -> tuple[str, list[ContextChunk]]:
    """Trim retrieval first, then memory, then project state until formatted context fits."""
    kept_chunks = list(chunks)
    formatted = _format_full_context(project_state_text, memory_text, kept_chunks)
    while len(formatted) > budget and kept_chunks:
        kept_chunks.pop()
        formatted = _format_full_context(project_state_text, memory_text, kept_chunks)

    if len(formatted) <= budget:
        return formatted, kept_chunks

    memory_text, _ = _truncate_text_block(memory_text, max(budget - len(project_state_text), 0))
    formatted = _format_full_context(project_state_text, memory_text, kept_chunks)
    if len(formatted) <= budget:
        return formatted, kept_chunks

    project_state_text, _ = _truncate_text_block(project_state_text, budget)
    formatted = _format_full_context(project_state_text, "", [])
    if len(formatted) > budget:
        formatted, _ = _truncate_text_block(formatted, budget)
    return formatted, []