headroom/learn/writer.py

"""Context writers — write learned patterns to agent-specific context files.

Writers take Recommendations and write them to the appropriate context
injection mechanism for each agent system (CLAUDE.md, .cursorrules, etc.).
"""

from __future__ import annotations

import re
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from pathlib import Path

from .models import (
    ProjectInfo,
    Recommendation,
    RecommendationTarget,
)

# Marker delimiters for Headroom-managed sections
_MARKER_START = "<!-- headroom:learn:start -->"
_MARKER_END = "<!-- headroom:learn:end -->"
_MARKER_PATTERN = re.compile(
    re.escape(_MARKER_START) + r".*?" + re.escape(_MARKER_END),
    re.DOTALL,
)


# =============================================================================
# Abstract Writer
# =============================================================================


class ContextWriter(ABC):
    """Base class for writing recommendations to context/memory files."""

    @abstractmethod
    def write(
        self,
        recommendations: list[Recommendation],
        project: ProjectInfo,
        dry_run: bool = True,
    ) -> WriteResult: ...


# =============================================================================
# Write Result
# =============================================================================


class WriteResult:
    """Result of a write operation."""

    def __init__(self) -> None:
        self.files_written: list[Path] = []
        self.content_by_file: dict[Path, str] = {}
        self.dry_run: bool = True

    def add(self, path: Path, content: str) -> None:
        self.files_written.append(path)
        self.content_by_file[path] = content


# =============================================================================
# Shared section builder
# =============================================================================


def _build_section(recommendations: list[Recommendation]) -> str:
    """Build the marker-delimited section content from recommendations."""
    now = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    lines = [
        _MARKER_START,
        "## Headroom Learned Patterns",
        f"*Auto-generated by `headroom learn` on {now} — do not edit manually*",
        "",
    ]

    for rec in recommendations:
        lines.append(f"### {rec.section}")
        if rec.estimated_tokens_saved > 0:
            lines.append(f"*~{rec.estimated_tokens_saved:,} tokens/session saved*")
        lines.append(rec.content)
        lines.append("")

    lines.append(_MARKER_END)
    return "\n".join(lines)


def _merge_into_file(file_path: Path, section: str) -> str:
    """Merge the section into an existing file, replacing any prior section."""
    if file_path.exists():
        existing = file_path.read_text()
        if _MARKER_START in existing:
            return _MARKER_PATTERN.sub(section, existing)
        return existing.rstrip() + "\n\n" + section + "\n"
    return section + "\n"


# =============================================================================
# Claude Code Writer
# =============================================================================


class ClaudeCodeWriter(ContextWriter):
    """Writes learned patterns to CLAUDE.md and MEMORY.md for Claude Code."""

    def write(
        self,
        recommendations: list[Recommendation],
        project: ProjectInfo,
        dry_run: bool = True,
    ) -> WriteResult:
        result = WriteResult()
        result.dry_run = dry_run

        context_recs = [r for r in recommendations if r.target == RecommendationTarget.CONTEXT_FILE]
        memory_recs = [r for r in recommendations if r.target == RecommendationTarget.MEMORY_FILE]

        if context_recs:
            claude_md_path = self._resolve_context_path(project)
            section_content = _build_section(context_recs)
            full_content = _merge_into_file(claude_md_path, section_content)
            result.add(claude_md_path, full_content)
            if not dry_run:
                claude_md_path.parent.mkdir(parents=True, exist_ok=True)
                claude_md_path.write_text(full_content)

        if memory_recs:
            memory_path = self._resolve_memory_path(project)
            section_content = _build_section(memory_recs)
            full_content = _merge_into_file(memory_path, section_content)
            result.add(memory_path, full_content)
            if not dry_run:
                memory_path.parent.mkdir(parents=True, exist_ok=True)
                memory_path.write_text(full_content)

        return result

    def _resolve_context_path(self, project: ProjectInfo) -> Path:
        if project.context_file:
            return project.context_file
        return project.project_path / "CLAUDE.md"

    def _resolve_memory_path(self, project: ProjectInfo) -> Path:
        if project.memory_file:
            return project.memory_file
        return project.data_path / "memory" / "MEMORY.md"


# =============================================================================
# Codex Writer (OpenAI Codex CLI)
# =============================================================================


class CodexWriter(ContextWriter):
    """Writes learned patterns to AGENTS.md and instructions.md for Codex CLI."""

    def write(
        self,
        recommendations: list[Recommendation],
        project: ProjectInfo,
        dry_run: bool = True,
    ) -> WriteResult:
        result = WriteResult()
        result.dry_run = dry_run

        context_recs = [r for r in recommendations if r.target == RecommendationTarget.CONTEXT_FILE]
        memory_recs = [r for r in recommendations if r.target == RecommendationTarget.MEMORY_FILE]

        if context_recs:
            agents_md = project.context_file or (project.project_path / "AGENTS.md")
            section_content = _build_section(context_recs)
            full_content = _merge_into_file(agents_md, section_content)
            result.add(agents_md, full_content)
            if not dry_run:
                agents_md.parent.mkdir(parents=True, exist_ok=True)
                agents_md.write_text(full_content)

        if memory_recs:
            instructions_md = project.memory_file or (project.data_path.parent / "instructions.md")
            section_content = _build_section(memory_recs)
            full_content = _merge_into_file(instructions_md, section_content)
            result.add(instructions_md, full_content)
            if not dry_run:
                instructions_md.parent.mkdir(parents=True, exist_ok=True)
                instructions_md.write_text(full_content)

        return result


# =============================================================================
# Gemini Writer (Google Gemini CLI)
# =============================================================================


class GeminiWriter(ContextWriter):
    """Writes learned patterns to GEMINI.md for Gemini CLI."""

    def write(
        self,
        recommendations: list[Recommendation],
        project: ProjectInfo,
        dry_run: bool = True,
    ) -> WriteResult:
        result = WriteResult()
        result.dry_run = dry_run

        if not recommendations:
            return result

        gemini_md = project.context_file or (project.project_path / "GEMINI.md")
        section_content = _build_section(recommendations)
        full_content = _merge_into_file(gemini_md, section_content)
        result.add(gemini_md, full_content)
        if not dry_run:
            gemini_md.parent.mkdir(parents=True, exist_ok=True)
            gemini_md.write_text(full_content)

        return result