/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 ._shared import claude_config_dir
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,
)


def _read_text_tolerant(file_path: Path) -> str:
    """Read an existing context file that we are about to rewrite as UTF-8.

    These files are predominantly valid UTF-8 but may carry a stray legacy
    byte (e.g. a cp1252 em-dash ``0x97``). Strict UTF-8 decoding aborts the
    whole ``--apply`` on a single such byte, so fall back to UTF-8 with
    replacement: this preserves the valid UTF-8 content — a full-file cp1252
    fallback would instead turn every genuine UTF-8 em-dash into mojibake —
    and the subsequent ``write_text(encoding="utf-8")`` self-heals the file.
    """
    raw = file_path.read_bytes()
    try:
        return raw.decode("utf-8")
    except UnicodeDecodeError:
        return raw.decode("utf-8", errors="replace")


# =============================================================================
# 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
        # Human-readable notices (e.g. legacy CLAUDE.md migration) surfaced by the CLI.
        self.warnings: list[str] = []

    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)


# Matches the "*~N tokens/session saved*" annotation emitted by _build_section.
_TOKENS_ANNOTATION_PATTERN = re.compile(r"\*~([\d,]+) tokens/session saved\*\n?")


def extract_marker_block(file_content: str) -> str | None:
    """Return the raw text of the headroom:learn marker block, or None.

    Unlike _parse_prior_recommendations, this returns the block verbatim
    (including the start/end markers) so it can be fed back to an LLM as
    context without losing formatting. Returns None if no block is present.
    """
    match = _MARKER_PATTERN.search(file_content)
    return match.group(0) if match else None


def _parse_prior_recommendations(existing: str) -> list[Recommendation]:
    """Parse recommendations out of a prior marker block.

    Returns [] if no marker block is present or it contains no sections.
    The returned Recommendation objects are round-trip compatible with
    _build_section — target is set to a placeholder since the marker block
    itself doesn't record it (blocks are always per-file and per-target).
    """
    match = _MARKER_PATTERN.search(existing)
    if not match:
        return []
    inner = match.group(0)[len(_MARKER_START) : -len(_MARKER_END)]

    recs: list[Recommendation] = []
    for part in re.split(r"\n### ", "\n" + inner)[1:]:
        heading_line, _, body = part.partition("\n")
        heading = heading_line.strip()
        if not heading:
            continue

        tokens_saved = 0
        tokens_match = _TOKENS_ANNOTATION_PATTERN.match(body)
        if tokens_match:
            tokens_saved = int(tokens_match.group(1).replace(",", ""))
            body = body[tokens_match.end() :]

        recs.append(
            Recommendation(
                target=RecommendationTarget.CONTEXT_FILE,
                section=heading,
                content=body.rstrip(),
                estimated_tokens_saved=tokens_saved,
            )
        )
    return recs


def _merge_recommendations(
    file_path: Path,
    new_recommendations: list[Recommendation],
) -> list[Recommendation]:
    """Union new recommendations with prior ones whose section is not re-surfaced.

    Sections produced by the current run take precedence over same-named
    prior sections — the latest analysis is authoritative. Prior sections
    whose headings do not reappear in the new run are carried forward so
    a re-run doesn't silently drop accumulated learnings. To fully rebuild
    the block, delete it manually and re-run.
    """
    if not file_path.exists():
        return new_recommendations
    prior = _parse_prior_recommendations(_read_text_tolerant(file_path))
    if not prior:
        return new_recommendations
    new_sections = {r.section for r in new_recommendations}
    carried = [p for p in prior if p.section not in new_sections]
    return list(new_recommendations) + carried


def _merge_into_file(file_path: Path, new_recommendations: list[Recommendation]) -> str:
    """Merge new recommendations with any existing marker block and rebuild the file."""
    merged = _merge_recommendations(file_path, new_recommendations)
    section = _build_section(merged)
    if file_path.exists():
        existing = _read_text_tolerant(file_path)
        if _MARKER_START in existing:
            return _MARKER_PATTERN.sub(lambda _match: section, existing)
        return existing.rstrip() + "\n\n" + section + "\n"
    return section + "\n"


def _strip_marker_block(content: str) -> str:
    """Remove the headroom:learn marker block from text, tidying blank lines.

    Used when migrating a stale block out of the team-shared CLAUDE.md into the
    personal CLAUDE.local.md. Returns "" if nothing but the block remained.
    """
    cleaned = _MARKER_PATTERN.sub("", content)
    cleaned = re.sub(r"\n{3,}", "\n\n", cleaned).strip()
    return cleaned + "\n" if cleaned else ""


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


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

    Project-level learnings default to ``CLAUDE.local.md`` rather than
    ``CLAUDE.md``: per Claude Code's memory convention ``CLAUDE.md`` is
    team-shared and checked into git, while ``CLAUDE.local.md`` is personal and
    gitignored by default. Learned patterns are personal-by-default (they hold
    machine-specific absolute paths and tool-discovery byproducts), so writing
    them to the shared file pollutes it for teammates (issue #1072).

    Pass an explicit target via :meth:`set_context_target` (CLI ``--target``) to
    override -- e.g. ``CLAUDE.md`` to opt back into the shared file.
    """

    def __init__(self, context_target: str | None = None) -> None:
        # Explicit write target for CONTEXT_FILE recs (overrides the default).
        self._context_target = context_target

    def set_context_target(self, context_target: str | None) -> None:
        """Override where CONTEXT_FILE recommendations are written.

        Accepts a path relative to the project root or an absolute path.
        """
        self._context_target = context_target

    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:
            target_path = self._resolve_context_path(project)
            # Migrate any stale block left in the team-shared CLAUDE.md by older
            # headroom versions into the new target, then strip it from CLAUDE.md
            # so the shared file is no longer polluted.
            migrated = self._migrate_legacy_block(project, target_path, result, dry_run)
            new_sections = {r.section for r in context_recs}
            merged_recs = context_recs + [r for r in migrated if r.section not in new_sections]
            full_content = _merge_into_file(target_path, merged_recs)
            result.add(target_path, full_content)
            if not dry_run:
                target_path.parent.mkdir(parents=True, exist_ok=True)
                target_path.write_text(full_content, encoding="utf-8")

        if memory_recs:
            memory_path = self._resolve_memory_path(project)
            full_content = _merge_into_file(memory_path, memory_recs)
            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, encoding="utf-8")

        return result

    def _resolve_context_path(self, project: ProjectInfo) -> Path:
        # Explicit --target wins over every default.
        if self._context_target is not None:
            target = Path(self._context_target).expanduser()
            return target if target.is_absolute() else project.project_path / target
        # The home directory's CLAUDE.md (~/.claude/CLAUDE.md, or
        # $CLAUDE_CONFIG_DIR/CLAUDE.md) is the user's personal global memory,
        # not a team-shared file, so keep writing there.
        if project.project_path == Path.home():
            return claude_config_dir() / "CLAUDE.md"
        # Project level: default to the gitignored, personal CLAUDE.local.md so
        # we never pollute the team-shared CLAUDE.md (issue #1072).
        return project.project_path / "CLAUDE.local.md"

    def _migrate_legacy_block(
        self,
        project: ProjectInfo,
        target_path: Path,
        result: WriteResult,
        dry_run: bool,
    ) -> list[Recommendation]:
        """Move a stale headroom block out of CLAUDE.md into the new target.

        Only fires for the default project-level case (no explicit --target, not
        the home directory) when CLAUDE.md still carries a marker block and the
        new target doesn't yet own one. Returns the migrated recommendations so
        the caller can carry them forward; records the cleaned CLAUDE.md and a
        warning on ``result``. Honors ``dry_run`` (no writes, warning still set).
        """
        legacy_path = project.project_path / "CLAUDE.md"
        if self._context_target is not None or project.project_path == Path.home():
            return []
        if target_path == legacy_path or not legacy_path.exists():
            return []
        legacy_text = _read_text_tolerant(legacy_path)
        if _MARKER_START not in legacy_text:
            return []
        # If the target already owns a block, it is the source of truth -- don't
        # double-migrate or clobber accumulated learnings.
        if target_path.exists() and _MARKER_START in _read_text_tolerant(target_path):
            return []

        migrated = _parse_prior_recommendations(legacy_text)
        cleaned = _strip_marker_block(legacy_text)
        gitignore_hint = f" Ensure {target_path.name} is in your .gitignore so it stays personal."
        if cleaned:
            # CLAUDE.md has hand-written content too — keep it, drop only the block.
            result.add(legacy_path, cleaned)
            result.warnings.append(
                f"Moved Headroom learnings out of {legacy_path} into {target_path}: "
                f"CLAUDE.md is team-shared, so personal learnings now live in "
                f"{target_path.name}. Review the diff before committing.{gitignore_hint}"
            )
            if not dry_run:
                legacy_path.write_text(cleaned, encoding="utf-8")
        else:
            # CLAUDE.md held nothing but the Headroom block — remove the husk.
            result.warnings.append(
                f"Removed {legacy_path} (it contained only Headroom learnings) and "
                f"moved them into {target_path}.{gitignore_hint}"
            )
            if not dry_run:
                legacy_path.unlink()
        return migrated

    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")
            full_content = _merge_into_file(agents_md, context_recs)
            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, encoding="utf-8")

        if memory_recs:
            instructions_md = project.memory_file or (project.data_path.parent / "instructions.md")
            full_content = _merge_into_file(instructions_md, memory_recs)
            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, encoding="utf-8")

        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")
        full_content = _merge_into_file(gemini_md, recommendations)
        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, encoding="utf-8")

        return result