narrate.py

"""Map-reduce narration: per-file summaries, then the whole-project story.

MAP   — summarise each file (style-neutral, cached by file path + content hash).
REDUCE— synthesise the project story in the chosen style + difficulty.

Both steps are grounded by `story.py` prompts. If the model endpoint is missing
or errors, we degrade to the deterministic `plain_fallback_story` so the app
still shows something true.
"""
from __future__ import annotations

import hashlib
from typing import Callable

import llm
import story
from schema import FileInfo, FileSummary, ProjectModel, ProjectStory

# Cache MAP summaries across runs/styles within a session: {content_hash: FileSummary}.
_SUMMARY_CACHE: dict[str, FileSummary] = {}

ProgressFn = Callable[[int, int, str], None]


def _hash(f: FileInfo) -> str:
    h = hashlib.sha1()  # noqa: S324 - cache key, not security
    h.update(f.path.encode())
    h.update(str([(s.kind, s.name) for s in f.symbols]).encode())
    h.update(str(sorted(f.depends_on)).encode())
    return h.hexdigest()


def summarise_files(model: ProjectModel, progress: ProgressFn | None = None
                    ) -> dict[str, FileSummary]:
    """MAP step. Returns {path: FileSummary}. Never raises; falls back per file."""
    out: dict[str, FileSummary] = {}
    total = len(model.files)
    use_model = llm.available()
    for i, f in enumerate(model.files, 1):
        if progress:
            progress(i, total, f.path)
        key = _hash(f)
        if key in _SUMMARY_CACHE:
            out[f.path] = _SUMMARY_CACHE[key]
            continue
        summary = _summarise_one(model, f) if use_model else _fallback_summary(f)
        _SUMMARY_CACHE[key] = summary
        out[f.path] = summary
    return out


def _summarise_one(model: ProjectModel, f: FileInfo) -> FileSummary:
    try:
        data = llm.chat_json(story.map_prompt(model, f),
                             schema=_file_schema(), max_tokens=1024)
        return FileSummary(path=f.path,
                           one_liner=(data.get("one_liner") or "").strip()[:80]
                           or _fallback_summary(f).one_liner,
                           summary=(data.get("summary") or "").strip()
                           or _fallback_summary(f).summary)
    except Exception:  # noqa: BLE001 - any backend hiccup -> safe fallback
        return _fallback_summary(f)


def _fallback_summary(f: FileInfo) -> FileSummary:
    import config
    role = config.ROLE_LABEL.get(f.role, f.role)
    names = ", ".join(s.name for s in f.symbols[:4])
    extra = f" It defines {names}." if names else ""
    return FileSummary(
        path=f.path, one_liner=role.lower(),
        summary=(f"This is part of the project's {role.lower()}.{extra} "
                 f"{f.fan_in} other file(s) depend on it."))


def tell_story(model: ProjectModel, summaries: dict[str, FileSummary],
               style_key: str, difficulty_key: str) -> ProjectStory:
    """REDUCE step. Falls back to the deterministic story on any failure."""
    if not llm.available():
        return story.plain_fallback_story(model, summaries)
    try:
        data = llm.chat_json(
            story.reduce_prompt(model, summaries, style_key, difficulty_key),
            schema=_story_schema(), max_tokens=2048)
        if not data.get("title"):
            raise ValueError("empty story")
        from schema import StorySection
        steps = [StorySection(heading=s.get("heading", ""), body=s.get("body", ""))
                 for s in data.get("steps", []) if s.get("body")]
        return ProjectStory(title=data["title"], overview=data.get("overview", ""),
                            steps=steps,
                            plain_overview=data.get("plain_overview", "") or data.get("overview", ""))
    except Exception:  # noqa: BLE001
        return story.plain_fallback_story(model, summaries)


def _file_schema() -> dict:
    from schema import file_summary_schema
    return file_summary_schema()


def _story_schema() -> dict:
    from schema import project_story_schema
    return project_story_schema()