Context Engine

#!/usr/bin/env python3
"""Context Analyzer - Analyze files and prompts for token usage, relevance scoring, and optimization.

Estimates token counts, scores context relevance across segments, identifies waste,
and provides actionable optimization suggestions based on the Context Engine's
Token Budget Allocation Framework.

Usage:
    python context_analyzer.py path/to/file_or_dir
    python context_analyzer.py --prompt "Fix the auth middleware bug"
    python context_analyzer.py path/to/dir --budget 128000 --json
"""

import argparse
import json
import math
import os
import re
import sys
from collections import Counter
from pathlib import Path


# Approximate token estimation: ~4 chars per token for English/code (GPT-family heuristic)
CHARS_PER_TOKEN = 4

# Token Budget Allocation Framework (from SKILL.md)
BUDGET_SEGMENTS = {
    "system_instructions": {"min": 0.05, "max": 0.10, "priority": "fixed"},
    "task_context": {"min": 0.20, "max": 0.30, "priority": "high"},
    "relevant_code": {"min": 0.25, "max": 0.40, "priority": "dynamic"},
    "conversation_history": {"min": 0.10, "max": 0.20, "priority": "sliding"},
    "tool_results": {"min": 0.05, "max": 0.15, "priority": "ephemeral"},
    "reserved_buffer": {"min": 0.05, "max": 0.10, "priority": "protected"},
}

# Patterns indicating low-signal content
LOW_SIGNAL_PATTERNS = [
    (r"^\s*#.*$", "comment_lines"),
    (r"^\s*$", "blank_lines"),
    (r"^\s*(import|from)\s+", "import_lines"),
    (r"^\s*\}\s*$", "closing_braces"),
    (r"^\s*pass\s*$", "pass_statements"),
    (r"^\s*\.{3}\s*$", "ellipsis"),
]

# Patterns indicating high-signal content
HIGH_SIGNAL_PATTERNS = [
    (r"^\s*(def|async def)\s+\w+", "function_definitions"),
    (r"^\s*class\s+\w+", "class_definitions"),
    (r"^\s*(return|yield)\s+", "return_statements"),
    (r"(raise|except|try|finally)\s+", "error_handling"),
    (r"^\s*@\w+", "decorators"),
]

VERBOSE_PATTERNS = [
    (r"#{3,}\s*-+\s*$", "decorative_dividers"),
    (r"^\s*#\s*={3,}", "decorative_headers"),
    (r"^\s*\"\"\"[\s\S]{200,}\"\"\"", "long_docstrings"),
    (r"(TODO|FIXME|HACK|XXX|NOQA)", "todo_markers"),
]


def estimate_tokens(text):
    """Estimate token count from text using character-based heuristic."""
    return max(1, len(text) // CHARS_PER_TOKEN)


def analyze_file(file_path):
    """Analyze a single file for token usage and content signals."""
    try:
        with open(file_path, "r", encoding="utf-8", errors="replace") as f:
            content = f.read()
    except (OSError, IOError) as e:
        return {"path": str(file_path), "error": str(e)}

    lines = content.splitlines()
    total_tokens = estimate_tokens(content)
    total_lines = len(lines)

    low_signal_counts = Counter()
    high_signal_counts = Counter()
    low_signal_lines = 0
    high_signal_lines = 0

    for line in lines:
        matched_low = False
        for pattern, name in LOW_SIGNAL_PATTERNS:
            if re.match(pattern, line):
                low_signal_counts[name] += 1
                low_signal_lines += 1
                matched_low = True
                break
        if not matched_low:
            for pattern, name in HIGH_SIGNAL_PATTERNS:
                if re.search(pattern, line):
                    high_signal_counts[name] += 1
                    high_signal_lines += 1
                    break

    verbose_counts = Counter()
    for pattern, name in VERBOSE_PATTERNS:
        matches = re.findall(pattern, content, re.MULTILINE)
        if matches:
            verbose_counts[name] = len(matches)

    neutral_lines = total_lines - low_signal_lines - high_signal_lines
    if total_lines > 0:
        relevance_score = round(
            (high_signal_lines * 1.0 + neutral_lines * 0.5) / total_lines, 3
        )
    else:
        relevance_score = 0.0

    # Detect duplicate/redundant blocks (repeated sequences of 3+ identical lines)
    redundant_blocks = 0
    if total_lines > 6:
        seen_blocks = set()
        for i in range(len(lines) - 2):
            block = tuple(lines[i : i + 3])
            if block in seen_blocks and any(l.strip() for l in block):
                redundant_blocks += 1
            seen_blocks.add(block)

    ext = Path(file_path).suffix.lower()
    size_bytes = os.path.getsize(file_path)

    return {
        "path": str(file_path),
        "extension": ext,
        "size_bytes": size_bytes,
        "total_lines": total_lines,
        "estimated_tokens": total_tokens,
        "relevance_score": relevance_score,
        "high_signal_lines": high_signal_lines,
        "low_signal_lines": low_signal_lines,
        "neutral_lines": neutral_lines,
        "high_signal_breakdown": dict(high_signal_counts),
        "low_signal_breakdown": dict(low_signal_counts),
        "verbose_patterns": dict(verbose_counts),
        "redundant_blocks": redundant_blocks,
    }


def compute_budget_analysis(total_tokens, budget):
    """Compute how the analyzed tokens fit within a context budget."""
    segments = {}
    for name, config in BUDGET_SEGMENTS.items():
        min_tokens = int(budget * config["min"])
        max_tokens = int(budget * config["max"])
        segments[name] = {
            "min_tokens": min_tokens,
            "max_tokens": max_tokens,
            "priority": config["priority"],
        }

    utilization = round(total_tokens / budget, 3) if budget > 0 else 0.0
    remaining = max(0, budget - total_tokens)

    return {
        "budget_tokens": budget,
        "used_tokens": total_tokens,
        "remaining_tokens": remaining,
        "utilization": utilization,
        "segments": segments,
        "over_budget": total_tokens > budget,
    }


def generate_suggestions(file_results, budget_analysis):
    """Generate optimization suggestions from analysis results."""
    suggestions = []

    total_low = sum(r.get("low_signal_lines", 0) for r in file_results if "error" not in r)
    total_lines = sum(r.get("total_lines", 0) for r in file_results if "error" not in r)

    if total_lines > 0 and total_low / total_lines > 0.35:
        suggestions.append({
            "type": "reduce_low_signal",
            "severity": "high",
            "message": (
                f"{total_low}/{total_lines} lines ({total_low * 100 // total_lines}%) are low-signal "
                "(blanks, comments, closing braces). Consider pruning or summarizing."
            ),
        })

    total_redundant = sum(r.get("redundant_blocks", 0) for r in file_results if "error" not in r)
    if total_redundant > 3:
        suggestions.append({
            "type": "remove_redundancy",
            "severity": "medium",
            "message": f"Detected {total_redundant} redundant 3-line blocks. Deduplicate repeated code sections.",
        })

    if budget_analysis and budget_analysis["over_budget"]:
        overage = budget_analysis["used_tokens"] - budget_analysis["budget_tokens"]
        suggestions.append({
            "type": "over_budget",
            "severity": "critical",
            "message": (
                f"Context exceeds budget by {overage:,} tokens. "
                "Apply tiered loading: keep Tier 0/1, summarize Tier 2, drop Tier 3."
            ),
        })
    elif budget_analysis and budget_analysis["utilization"] < 0.5:
        suggestions.append({
            "type": "under_utilized",
            "severity": "info",
            "message": (
                f"Only {budget_analysis['utilization'] * 100:.0f}% of budget used. "
                "You can load additional dependency or test files for better coverage."
            ),
        })

    # Check for very large individual files
    for r in file_results:
        if "error" in r:
            continue
        if r["estimated_tokens"] > 8000:
            suggestions.append({
                "type": "large_file",
                "severity": "medium",
                "message": (
                    f"{r['path']} is ~{r['estimated_tokens']:,} tokens. "
                    "Consider loading only relevant functions instead of the full file."
                ),
            })

    if not suggestions:
        suggestions.append({
            "type": "ok",
            "severity": "info",
            "message": "Context looks well-optimized. No major issues detected.",
        })

    return suggestions


def collect_files(path, max_files=100):
    """Collect files from a path (file or directory)."""
    p = Path(path)
    if p.is_file():
        return [p]

    if p.is_dir():
        files = []
        skip_dirs = {".git", "__pycache__", "node_modules", ".venv", "venv", ".tox"}
        for root, dirs, filenames in os.walk(p):
            dirs[:] = [d for d in dirs if d not in skip_dirs]
            for fname in filenames:
                fp = Path(root) / fname
                if fp.suffix.lower() in {
                    ".py", ".js", ".ts", ".tsx", ".jsx", ".md", ".yaml", ".yml",
                    ".json", ".toml", ".cfg", ".ini", ".sh", ".go", ".rs", ".java",
                    ".rb", ".c", ".cpp", ".h", ".hpp", ".cs", ".swift", ".kt",
                }:
                    files.append(fp)
                    if len(files) >= max_files:
                        return files
        return files

    return []


def format_human(file_results, budget_analysis, suggestions):
    """Format results for human-readable terminal output."""
    lines = []
    lines.append("=" * 64)
    lines.append("  CONTEXT ANALYZER REPORT")
    lines.append("=" * 64)

    total_tokens = sum(r.get("estimated_tokens", 0) for r in file_results if "error" not in r)
    total_files = sum(1 for r in file_results if "error" not in r)
    lines.append(f"\n  Files analyzed: {total_files}")
    lines.append(f"  Total estimated tokens: {total_tokens:,}")

    if budget_analysis:
        lines.append(f"  Budget: {budget_analysis['budget_tokens']:,} tokens")
        lines.append(f"  Utilization: {budget_analysis['utilization'] * 100:.1f}%")
        if budget_analysis["over_budget"]:
            lines.append(f"  ** OVER BUDGET by {total_tokens - budget_analysis['budget_tokens']:,} tokens **")

    lines.append(f"\n{'  File':<45} {'Tokens':>8} {'Relevance':>10}")
    lines.append("  " + "-" * 62)
    for r in sorted(file_results, key=lambda x: x.get("estimated_tokens", 0), reverse=True):
        if "error" in r:
            lines.append(f"  {r['path']:<45} ERROR: {r['error']}")
            continue
        name = r["path"]
        if len(name) > 43:
            name = "..." + name[-40:]
        lines.append(f"  {name:<45} {r['estimated_tokens']:>8,} {r['relevance_score']:>10.2f}")

    lines.append(f"\n  OPTIMIZATION SUGGESTIONS")
    lines.append("  " + "-" * 40)
    for s in suggestions:
        marker = {"critical": "!!!", "high": "!!", "medium": "!", "info": " "}
        icon = marker.get(s["severity"], " ")
        lines.append(f"  [{icon}] {s['message']}")

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


def main():
    parser = argparse.ArgumentParser(
        description="Analyze files/prompts for token usage, context relevance, and optimization suggestions.",
        epilog="Example: python context_analyzer.py src/ --budget 128000 --json",
    )
    parser.add_argument("path", nargs="?", help="File or directory to analyze")
    parser.add_argument("--prompt", type=str, help="Analyze a prompt string for token estimation")
    parser.add_argument("--budget", type=int, default=128000, help="Context window budget in tokens (default: 128000)")
    parser.add_argument("--max-files", type=int, default=100, help="Max files to analyze in a directory (default: 100)")
    parser.add_argument("--json", action="store_true", dest="json_output", help="Output results as JSON")

    args = parser.parse_args()

    if not args.path and not args.prompt:
        parser.print_help()
        sys.exit(1)

    file_results = []

    if args.prompt:
        tokens = estimate_tokens(args.prompt)
        file_results.append({
            "path": "<prompt>",
            "extension": "",
            "size_bytes": len(args.prompt.encode("utf-8")),
            "total_lines": args.prompt.count("\n") + 1,
            "estimated_tokens": tokens,
            "relevance_score": 1.0,
            "high_signal_lines": args.prompt.count("\n") + 1,
            "low_signal_lines": 0,
            "neutral_lines": 0,
            "high_signal_breakdown": {},
            "low_signal_breakdown": {},
            "verbose_patterns": {},
            "redundant_blocks": 0,
        })

    if args.path:
        files = collect_files(args.path, max_files=args.max_files)
        if not files:
            print(f"Error: No analyzable files found at '{args.path}'", file=sys.stderr)
            sys.exit(1)
        for f in files:
            file_results.append(analyze_file(f))

    total_tokens = sum(r.get("estimated_tokens", 0) for r in file_results if "error" not in r)
    budget_analysis = compute_budget_analysis(total_tokens, args.budget)
    suggestions = generate_suggestions(file_results, budget_analysis)

    result = {
        "summary": {
            "files_analyzed": sum(1 for r in file_results if "error" not in r),
            "total_estimated_tokens": total_tokens,
            "budget": args.budget,
        },
        "files": file_results,
        "budget_analysis": budget_analysis,
        "suggestions": suggestions,
    }

    if args.json_output:
        print(json.dumps(result, indent=2))
    else:
        print(format_human(file_results, budget_analysis, suggestions))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Context Pruner - Intelligently prune context by removing low-relevance content.

Removes blank lines, comment blocks, redundant sections, verbose patterns,
and boilerplate to compress context while preserving high-signal code.
Implements the Adaptive Compression pattern from the Context Engine skill.

Usage:
    python context_pruner.py path/to/file.py
    python context_pruner.py path/to/file.py --aggressive --json
    python context_pruner.py path/to/dir --output pruned_output/
    cat context.txt | python context_pruner.py --stdin
"""

import argparse
import json
import os
import re
import sys
from pathlib import Path


# Pruning levels control how aggressively we strip content
PRUNING_LEVELS = {
    "light": {
        "strip_blank_runs": 2,       # Collapse runs of blanks to max N
        "strip_comments": False,     # Keep all comments
        "strip_docstrings": False,   # Keep all docstrings
        "strip_imports": False,      # Keep all imports
        "strip_decorative": True,    # Remove decorative dividers
        "collapse_braces": False,    # Keep closing brace lines
    },
    "moderate": {
        "strip_blank_runs": 1,
        "strip_comments": True,      # Remove standalone comment lines
        "strip_docstrings": False,
        "strip_imports": False,
        "strip_decorative": True,
        "collapse_braces": True,
    },
    "aggressive": {
        "strip_blank_runs": 0,       # Remove ALL blank lines
        "strip_comments": True,
        "strip_docstrings": True,    # Trim long docstrings to first line
        "strip_imports": True,       # Collapse import blocks to summary
        "strip_decorative": True,
        "collapse_braces": True,
    },
}


def collapse_blank_runs(lines, max_consecutive):
    """Collapse consecutive blank lines to at most max_consecutive."""
    result = []
    blank_count = 0
    for line in lines:
        if line.strip() == "":
            blank_count += 1
            if blank_count <= max_consecutive:
                result.append(line)
        else:
            blank_count = 0
            result.append(line)
    return result


def strip_comment_lines(lines):
    """Remove standalone comment lines (preserving inline comments and shebangs)."""
    result = []
    for i, line in enumerate(lines):
        stripped = line.strip()
        # Keep shebangs, keep inline comments (code before #)
        if stripped.startswith("#") and not stripped.startswith("#!"):
            # Keep type: ignore, noqa, pragma, and encoding comments
            if any(kw in stripped.lower() for kw in ["type:", "noqa", "pragma", "coding", "pylint", "fmt:"]):
                result.append(line)
                continue
            continue
        # Strip // comments for JS/TS/Go/Java/C etc.
        if stripped.startswith("//") and not stripped.startswith("///"):
            continue
        result.append(line)
    return result


def trim_docstrings(lines):
    """Trim multi-line docstrings to just the first summary line."""
    result = []
    in_docstring = False
    docstring_quote = None
    docstring_indent = ""
    first_line_added = False

    for line in lines:
        stripped = line.strip()

        if not in_docstring:
            # Detect docstring opening
            for quote in ['"""', "'''"]:
                if quote in stripped:
                    idx = stripped.index(quote)
                    after_quote = stripped[idx + 3:]
                    # Single-line docstring — keep as is
                    if quote in after_quote:
                        result.append(line)
                        break
                    # Multi-line docstring starts
                    in_docstring = True
                    docstring_quote = quote
                    docstring_indent = line[: len(line) - len(line.lstrip())]
                    # Keep the opening line
                    result.append(line)
                    first_line_added = True
                    break
            else:
                result.append(line)
        else:
            # Inside a docstring — look for closing quote
            if docstring_quote in stripped:
                # Close the docstring with just the closing quotes
                result.append(docstring_indent + docstring_quote)
                in_docstring = False
                docstring_quote = None
                first_line_added = False
            # Skip intermediate docstring lines (they are pruned)

    return result


def collapse_imports(lines):
    """Collapse consecutive import lines into a summary comment."""
    result = []
    import_block = []
    import_modules = set()

    def flush_imports():
        if import_block:
            count = len(import_block)
            # Extract module names
            for imp_line in import_block:
                m = re.match(r"^\s*(?:from\s+(\S+)|import\s+(\S+))", imp_line)
                if m:
                    import_modules.add(m.group(1) or m.group(2))
            modules_str = ", ".join(sorted(import_modules)[:8])
            if len(import_modules) > 8:
                modules_str += f", ... (+{len(import_modules) - 8} more)"
            result.append(f"# [{count} imports: {modules_str}]")
            import_block.clear()
            import_modules.clear()

    for line in lines:
        stripped = line.strip()
        if re.match(r"^(import |from \S+ import )", stripped):
            import_block.append(line)
        else:
            flush_imports()
            result.append(line)

    flush_imports()
    return result


def strip_decorative(lines):
    """Remove decorative dividers and banner comments."""
    result = []
    for line in lines:
        stripped = line.strip()
        # Lines that are purely decorative: ####, -----, ====, /****/
        if re.match(r"^[#\-=*\/\\]{4,}\s*$", stripped):
            continue
        # Banner-style comments like # ========== SECTION ==========
        if re.match(r"^#\s*[=\-*]{3,}.*[=\-*]{3,}\s*$", stripped):
            # Keep the text, strip the decoration
            text = re.sub(r"[=\-*]+", "", stripped.lstrip("#")).strip()
            if text:
                indent = line[: len(line) - len(line.lstrip())]
                result.append(f"{indent}# {text}")
            continue
        result.append(line)
    return result


def collapse_closing_braces(lines):
    """Remove lines that are only closing braces/brackets with optional whitespace."""
    result = []
    for line in lines:
        stripped = line.strip()
        if stripped in ("}", "},", ");", "});", "]", "],"):
            # Keep it but strip trailing whitespace
            result.append(line.rstrip())
        else:
            result.append(line)
    # Remove consecutive closing-only lines (keep just one)
    final = []
    prev_closing = False
    for line in result:
        stripped = line.strip()
        is_closing = stripped in ("}", "},", ");", "});", "]", "],")
        if is_closing and prev_closing:
            # Merge onto previous line
            if final:
                final[-1] = final[-1] + " " + stripped
            continue
        final.append(line)
        prev_closing = is_closing
    return final


def detect_redundant_blocks(lines, min_block_size=3):
    """Detect and remove repeated blocks of identical lines."""
    result = []
    seen_blocks = {}
    i = 0
    removed = 0

    while i < len(lines):
        if i + min_block_size <= len(lines):
            block = tuple(l.strip() for l in lines[i : i + min_block_size])
            # Skip blocks that are all empty
            if any(l for l in block):
                if block in seen_blocks:
                    result.append(f"# [duplicate block removed, first seen line {seen_blocks[block] + 1}]")
                    i += min_block_size
                    removed += min_block_size
                    continue
                seen_blocks[block] = i
        result.append(lines[i])
        i += 1

    return result, removed


def prune_content(content, level="moderate"):
    """Apply pruning strategies to content string. Returns (pruned_content, stats)."""
    config = PRUNING_LEVELS.get(level, PRUNING_LEVELS["moderate"])
    lines = content.splitlines()
    original_lines = len(lines)
    original_tokens = max(1, len(content) // 4)
    stats = {"original_lines": original_lines, "original_tokens": original_tokens, "operations": []}

    # 1. Strip decorative elements
    if config["strip_decorative"]:
        before = len(lines)
        lines = strip_decorative(lines)
        removed = before - len(lines)
        if removed > 0:
            stats["operations"].append({"op": "strip_decorative", "lines_removed": removed})

    # 2. Strip comments
    if config["strip_comments"]:
        before = len(lines)
        lines = strip_comment_lines(lines)
        removed = before - len(lines)
        if removed > 0:
            stats["operations"].append({"op": "strip_comments", "lines_removed": removed})

    # 3. Trim docstrings
    if config["strip_docstrings"]:
        before = len(lines)
        lines = trim_docstrings(lines)
        removed = before - len(lines)
        if removed > 0:
            stats["operations"].append({"op": "trim_docstrings", "lines_removed": removed})

    # 4. Collapse imports
    if config["strip_imports"]:
        before = len(lines)
        lines = collapse_imports(lines)
        removed = before - len(lines)
        if removed > 0:
            stats["operations"].append({"op": "collapse_imports", "lines_removed": removed})

    # 5. Collapse closing braces
    if config["collapse_braces"]:
        before = len(lines)
        lines = collapse_closing_braces(lines)
        removed = before - len(lines)
        if removed > 0:
            stats["operations"].append({"op": "collapse_braces", "lines_removed": removed})

    # 6. Remove redundant blocks
    before = len(lines)
    lines, dup_removed = detect_redundant_blocks(lines)
    if dup_removed > 0:
        stats["operations"].append({"op": "remove_duplicates", "lines_removed": dup_removed})

    # 7. Collapse blank lines (always, last step)
    before = len(lines)
    lines = collapse_blank_runs(lines, config["strip_blank_runs"])
    removed = before - len(lines)
    if removed > 0:
        stats["operations"].append({"op": "collapse_blanks", "lines_removed": removed})

    pruned_content = "\n".join(lines)
    pruned_tokens = max(1, len(pruned_content) // 4)

    stats["pruned_lines"] = len(lines)
    stats["pruned_tokens"] = pruned_tokens
    stats["lines_removed"] = original_lines - len(lines)
    stats["tokens_saved"] = original_tokens - pruned_tokens
    stats["compression_ratio"] = round(pruned_tokens / original_tokens, 3) if original_tokens > 0 else 1.0

    return pruned_content, stats


def prune_file(file_path, level, output_dir=None):
    """Prune a single file and optionally write output."""
    try:
        with open(file_path, "r", encoding="utf-8", errors="replace") as f:
            content = f.read()
    except (OSError, IOError) as e:
        return {"path": str(file_path), "error": str(e)}

    pruned, stats = prune_content(content, level)
    stats["path"] = str(file_path)

    if output_dir:
        out_path = Path(output_dir) / Path(file_path).name
        os.makedirs(output_dir, exist_ok=True)
        with open(out_path, "w", encoding="utf-8") as f:
            f.write(pruned)
        stats["output_path"] = str(out_path)

    stats["pruned_content"] = pruned
    return stats


def collect_files(path, max_files=50):
    """Collect prunable files from a path."""
    p = Path(path)
    if p.is_file():
        return [p]
    if p.is_dir():
        files = []
        skip = {".git", "__pycache__", "node_modules", ".venv", "venv"}
        exts = {".py", ".js", ".ts", ".tsx", ".jsx", ".md", ".yaml", ".yml",
                ".go", ".rs", ".java", ".rb", ".c", ".cpp", ".h", ".sh"}
        for root, dirs, fnames in os.walk(p):
            dirs[:] = [d for d in dirs if d not in skip]
            for fn in fnames:
                fp = Path(root) / fn
                if fp.suffix.lower() in exts:
                    files.append(fp)
                    if len(files) >= max_files:
                        return files
        return files
    return []


def format_human(results):
    """Format results for human-readable output."""
    lines = []
    lines.append("=" * 60)
    lines.append("  CONTEXT PRUNER REPORT")
    lines.append("=" * 60)

    total_saved = 0
    total_original = 0

    for r in results:
        if "error" in r:
            lines.append(f"\n  {r['path']}: ERROR - {r['error']}")
            continue

        lines.append(f"\n  File: {r['path']}")
        lines.append(f"    Lines: {r['original_lines']} -> {r['pruned_lines']} ({r['lines_removed']} removed)")
        lines.append(f"    Tokens: ~{r['original_tokens']:,} -> ~{r['pruned_tokens']:,} ({r['tokens_saved']:,} saved)")
        lines.append(f"    Compression: {r['compression_ratio'] * 100:.1f}%")

        if r.get("operations"):
            lines.append("    Operations:")
            for op in r["operations"]:
                lines.append(f"      - {op['op']}: {op['lines_removed']} lines removed")

        if "output_path" in r:
            lines.append(f"    Written to: {r['output_path']}")

        total_saved += r.get("tokens_saved", 0)
        total_original += r.get("original_tokens", 0)

    lines.append("\n" + "-" * 60)
    lines.append(f"  Total tokens saved: ~{total_saved:,}")
    if total_original > 0:
        lines.append(f"  Overall compression: {(total_original - total_saved) / total_original * 100:.1f}%")
    lines.append("")
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(
        description="Prune context by removing low-relevance content, redundancy, and verbose patterns.",
        epilog="Example: python context_pruner.py src/main.py --aggressive --json",
    )
    parser.add_argument("path", nargs="?", help="File or directory to prune")
    parser.add_argument("--stdin", action="store_true", help="Read content from stdin")
    parser.add_argument("--level", choices=["light", "moderate", "aggressive"], default="moderate",
                        help="Pruning aggressiveness (default: moderate)")
    parser.add_argument("--aggressive", action="store_true", help="Shortcut for --level aggressive")
    parser.add_argument("--output", type=str, help="Output directory for pruned files")
    parser.add_argument("--max-files", type=int, default=50, help="Max files to process (default: 50)")
    parser.add_argument("--json", action="store_true", dest="json_output", help="Output results as JSON")

    args = parser.parse_args()
    level = "aggressive" if args.aggressive else args.level

    if not args.path and not args.stdin:
        parser.print_help()
        sys.exit(1)

    results = []

    if args.stdin:
        content = sys.stdin.read()
        pruned, stats = prune_content(content, level)
        stats["path"] = "<stdin>"
        if not args.json_output:
            print(pruned)
            return
        stats.pop("pruned_content", None)
        results.append(stats)
    elif args.path:
        files = collect_files(args.path, max_files=args.max_files)
        if not files:
            print(f"Error: No files found at '{args.path}'", file=sys.stderr)
            sys.exit(1)
        for f in files:
            r = prune_file(f, level, output_dir=args.output)
            r.pop("pruned_content", None)
            results.append(r)

    if args.json_output:
        print(json.dumps(results, indent=2))
    else:
        print(format_human(results))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Memory Indexer - Index and search a memory/knowledge base directory.

Scans a directory of markdown/text files, builds a term-frequency index,
and scores entries by relevance to a query using TF-IDF-inspired ranking.
Designed for the Context Engine's Session Memory and Knowledge Base layers.

Usage:
    python memory_indexer.py path/to/knowledge_base --query "authentication flow"
    python memory_indexer.py path/to/memory_dir --query "deployment" --top 5 --json
    python memory_indexer.py path/to/dir --index-only --json
    python memory_indexer.py path/to/dir --stats
"""

import argparse
import json
import math
import os
import re
import sys
from collections import Counter, defaultdict
from datetime import datetime
from pathlib import Path


# Stop words to exclude from indexing
STOP_WORDS = frozenset({
    "a", "an", "the", "and", "or", "but", "in", "on", "at", "to", "for",
    "of", "with", "by", "from", "is", "are", "was", "were", "be", "been",
    "has", "have", "had", "do", "does", "did", "will", "would", "could",
    "should", "may", "might", "can", "this", "that", "these", "those",
    "it", "its", "not", "no", "if", "then", "else", "when", "which",
    "what", "how", "who", "where", "why", "all", "each", "every", "both",
    "as", "so", "up", "out", "about", "into", "over", "after", "before",
    "between", "under", "above", "such", "only", "also", "than", "too",
    "very", "just", "more", "most", "other", "some", "any", "new", "old",
    "use", "used", "using", "see", "e", "g", "i", "we", "you", "they",
})

# File extensions to index
INDEXABLE_EXTENSIONS = {".md", ".txt", ".yaml", ".yml", ".json", ".toml", ".rst", ".org"}

# Boost factors for different content locations
TITLE_BOOST = 3.0
HEADING_BOOST = 2.0
FRONTMATTER_BOOST = 1.5
CODE_BLOCK_PENALTY = 0.5


def tokenize(text):
    """Split text into lowercase tokens, filtering stop words and short tokens."""
    # Split on non-alphanumeric (keep underscores and hyphens for code terms)
    raw_tokens = re.findall(r"[a-zA-Z_][a-zA-Z0-9_-]*", text.lower())
    return [t for t in raw_tokens if t not in STOP_WORDS and len(t) > 1]


def extract_sections(content):
    """Extract structured sections from a markdown file."""
    sections = []
    current_heading = ""
    current_body = []
    in_code_block = False
    in_frontmatter = False
    frontmatter_text = []

    lines = content.splitlines()

    for i, line in enumerate(lines):
        stripped = line.strip()

        # Detect YAML frontmatter
        if i == 0 and stripped == "---":
            in_frontmatter = True
            continue
        if in_frontmatter:
            if stripped == "---":
                in_frontmatter = False
                sections.append({
                    "type": "frontmatter",
                    "heading": "_frontmatter",
                    "content": "\n".join(frontmatter_text),
                    "line": 0,
                })
                continue
            frontmatter_text.append(line)
            continue

        # Track code blocks
        if stripped.startswith("```"):
            in_code_block = not in_code_block
            current_body.append(line)
            continue

        # Detect headings
        if not in_code_block and re.match(r"^#{1,6}\s+", stripped):
            # Flush previous section
            if current_heading or current_body:
                sections.append({
                    "type": "section",
                    "heading": current_heading,
                    "content": "\n".join(current_body),
                    "line": max(0, i - len(current_body)),
                })
            current_heading = re.sub(r"^#{1,6}\s+", "", stripped)
            current_body = []
        else:
            current_body.append(line)

    # Flush last section
    if current_heading or current_body:
        sections.append({
            "type": "section",
            "heading": current_heading,
            "content": "\n".join(current_body),
            "line": max(0, len(lines) - len(current_body)),
        })

    return sections


def build_file_entry(file_path):
    """Build an index entry for a single file."""
    try:
        with open(file_path, "r", encoding="utf-8", errors="replace") as f:
            content = f.read()
    except (OSError, IOError) as e:
        return None

    stat = os.stat(file_path)
    sections = extract_sections(content)

    # Build term frequency map with positional boosts
    term_freq = Counter()
    section_terms = {}

    # Index file name tokens with title boost
    name_tokens = tokenize(Path(file_path).stem.replace("-", " ").replace("_", " "))
    for t in name_tokens:
        term_freq[t] += TITLE_BOOST

    for section in sections:
        heading_tokens = tokenize(section["heading"])
        body_tokens = tokenize(section["content"])

        # Boost heading terms
        for t in heading_tokens:
            boost = FRONTMATTER_BOOST if section["type"] == "frontmatter" else HEADING_BOOST
            term_freq[t] += boost

        # Body terms at base weight (penalize code blocks slightly)
        in_code = False
        for line in section["content"].splitlines():
            if line.strip().startswith("```"):
                in_code = not in_code
                continue
            line_tokens = tokenize(line)
            weight = CODE_BLOCK_PENALTY if in_code else 1.0
            for t in line_tokens:
                term_freq[t] += weight

        section_key = section["heading"] or f"_section_{section['line']}"
        section_terms[section_key] = set(heading_tokens + body_tokens)

    total_tokens = sum(term_freq.values())

    # Extract title from first heading or filename
    title = Path(file_path).stem
    for s in sections:
        if s["type"] == "section" and s["heading"]:
            title = s["heading"]
            break

    # Detect staleness based on modification time
    mtime = datetime.fromtimestamp(stat.st_mtime)
    days_old = (datetime.now() - mtime).days
    if days_old < 7:
        freshness = "fresh"
    elif days_old < 30:
        freshness = "aging"
    else:
        freshness = "stale"

    return {
        "path": str(file_path),
        "title": title,
        "size_bytes": stat.st_size,
        "modified": mtime.isoformat(),
        "days_since_modified": days_old,
        "freshness": freshness,
        "total_weighted_tokens": round(total_tokens, 1),
        "unique_terms": len(term_freq),
        "sections": [{"heading": s["heading"], "type": s["type"], "line": s["line"]}
                      for s in sections],
        "term_freq": dict(term_freq),
        "section_terms": {k: list(v) for k, v in section_terms.items()},
    }


def build_index(directory, max_files=500):
    """Build a full index of a knowledge base directory."""
    entries = []
    doc_freq = Counter()  # How many documents contain each term
    skip_dirs = {".git", "__pycache__", "node_modules", ".venv", "venv"}

    p = Path(directory)
    if not p.is_dir():
        return [], {}

    file_count = 0
    for root, dirs, fnames in os.walk(p):
        dirs[:] = [d for d in dirs if d not in skip_dirs]
        for fn in sorted(fnames):
            fp = Path(root) / fn
            if fp.suffix.lower() not in INDEXABLE_EXTENSIONS:
                continue
            entry = build_file_entry(fp)
            if entry:
                entries.append(entry)
                # Track document frequency
                for term in entry["term_freq"]:
                    doc_freq[term] += 1
                file_count += 1
                if file_count >= max_files:
                    break
        if file_count >= max_files:
            break

    return entries, dict(doc_freq)


def score_query(query, entries, doc_freq):
    """Score all entries against a query using TF-IDF-inspired ranking."""
    query_tokens = tokenize(query)
    if not query_tokens:
        return []

    num_docs = max(len(entries), 1)
    results = []

    for entry in entries:
        tf = entry["term_freq"]
        total = max(entry["total_weighted_tokens"], 1)
        score = 0.0
        matched_terms = []

        for qt in query_tokens:
            if qt in tf:
                # TF component: term frequency normalized by document length
                tf_val = tf[qt] / total
                # IDF component: inverse document frequency
                df = doc_freq.get(qt, 1)
                idf_val = math.log(num_docs / df) + 1
                term_score = tf_val * idf_val
                score += term_score
                matched_terms.append(qt)

        # Freshness bonus: fresh docs get a small boost
        if entry["freshness"] == "fresh":
            score *= 1.1
        elif entry["freshness"] == "stale":
            score *= 0.9

        # Coverage bonus: matching more query terms is better
        if len(query_tokens) > 1:
            coverage = len(set(matched_terms)) / len(set(query_tokens))
            score *= (0.5 + 0.5 * coverage)

        if score > 0:
            # Find which sections matched
            matched_sections = []
            query_set = set(query_tokens)
            for sec_name, sec_terms in entry.get("section_terms", {}).items():
                if query_set & set(sec_terms):
                    matched_sections.append(sec_name)

            results.append({
                "path": entry["path"],
                "title": entry["title"],
                "score": round(score, 4),
                "matched_terms": list(set(matched_terms)),
                "matched_sections": matched_sections,
                "freshness": entry["freshness"],
                "days_old": entry["days_since_modified"],
            })

    results.sort(key=lambda x: x["score"], reverse=True)
    return results


def compute_stats(entries):
    """Compute summary statistics for the index."""
    if not entries:
        return {"total_files": 0}

    total_terms = sum(e["unique_terms"] for e in entries)
    total_size = sum(e["size_bytes"] for e in entries)
    freshness_dist = Counter(e["freshness"] for e in entries)

    # Find most common terms across corpus
    global_freq = Counter()
    for e in entries:
        for term, freq in e["term_freq"].items():
            global_freq[term] += freq

    return {
        "total_files": len(entries),
        "total_size_bytes": total_size,
        "total_unique_terms": len(global_freq),
        "avg_terms_per_file": round(total_terms / len(entries), 1),
        "freshness_distribution": dict(freshness_dist),
        "top_terms": [{"term": t, "frequency": round(f, 1)} for t, f in global_freq.most_common(20)],
        "largest_files": sorted(
            [{"path": e["path"], "size": e["size_bytes"], "terms": e["unique_terms"]}
             for e in entries],
            key=lambda x: x["size"], reverse=True,
        )[:10],
    }


def format_search_results(results, query, top_n):
    """Format search results for human-readable output."""
    lines = []
    lines.append("=" * 60)
    lines.append(f"  MEMORY INDEX SEARCH: \"{query}\"")
    lines.append("=" * 60)

    if not results:
        lines.append("\n  No matching entries found.")
        lines.append("")
        return "\n".join(lines)

    shown = results[:top_n]
    lines.append(f"\n  Found {len(results)} matching entries (showing top {len(shown)}):\n")

    for i, r in enumerate(shown, 1):
        freshness_marker = {"fresh": "+", "aging": "~", "stale": "-"}.get(r["freshness"], "?")
        lines.append(f"  {i}. [{freshness_marker}] {r['title']}")
        lines.append(f"     Path: {r['path']}")
        lines.append(f"     Score: {r['score']:.4f}  |  Age: {r['days_old']}d ({r['freshness']})")
        lines.append(f"     Matched: {', '.join(r['matched_terms'])}")
        if r["matched_sections"]:
            lines.append(f"     Sections: {', '.join(r['matched_sections'][:5])}")
        lines.append("")

    lines.append(f"  Legend: [+] fresh (<7d)  [~] aging (7-30d)  [-] stale (>30d)")
    lines.append("")
    return "\n".join(lines)


def format_stats(stats):
    """Format index statistics for human-readable output."""
    lines = []
    lines.append("=" * 60)
    lines.append("  MEMORY INDEX STATISTICS")
    lines.append("=" * 60)

    lines.append(f"\n  Total files indexed: {stats['total_files']}")
    lines.append(f"  Total size: {stats.get('total_size_bytes', 0):,} bytes")
    lines.append(f"  Unique terms: {stats.get('total_unique_terms', 0):,}")
    lines.append(f"  Avg terms/file: {stats.get('avg_terms_per_file', 0)}")

    fd = stats.get("freshness_distribution", {})
    if fd:
        lines.append(f"\n  Freshness: {fd.get('fresh', 0)} fresh, {fd.get('aging', 0)} aging, {fd.get('stale', 0)} stale")

    top = stats.get("top_terms", [])
    if top:
        lines.append("\n  Top terms:")
        for t in top[:15]:
            lines.append(f"    {t['term']:<30} {t['frequency']:>8.0f}")

    largest = stats.get("largest_files", [])
    if largest:
        lines.append("\n  Largest files:")
        for f in largest[:5]:
            name = f["path"]
            if len(name) > 45:
                name = "..." + name[-42:]
            lines.append(f"    {name:<48} {f['size']:>8,} bytes")

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


def main():
    parser = argparse.ArgumentParser(
        description="Index and search a memory/knowledge base directory with TF-IDF relevance scoring.",
        epilog="Example: python memory_indexer.py docs/ --query 'auth middleware' --top 5",
    )
    parser.add_argument("directory", help="Directory containing knowledge base files to index")
    parser.add_argument("--query", "-q", type=str, help="Search query to score entries against")
    parser.add_argument("--top", "-n", type=int, default=10, help="Number of top results to return (default: 10)")
    parser.add_argument("--index-only", action="store_true", help="Build and display the index without searching")
    parser.add_argument("--stats", action="store_true", help="Show index statistics")
    parser.add_argument("--max-files", type=int, default=500, help="Maximum files to index (default: 500)")
    parser.add_argument("--json", action="store_true", dest="json_output", help="Output results as JSON")

    args = parser.parse_args()

    if not os.path.isdir(args.directory):
        print(f"Error: '{args.directory}' is not a directory", file=sys.stderr)
        sys.exit(1)

    if not args.query and not args.index_only and not args.stats:
        parser.print_help()
        print("\nError: Provide --query, --index-only, or --stats", file=sys.stderr)
        sys.exit(1)

    entries, doc_freq = build_index(args.directory, max_files=args.max_files)

    if not entries:
        print(f"No indexable files found in '{args.directory}'", file=sys.stderr)
        sys.exit(1)

    if args.stats:
        stats = compute_stats(entries)
        if args.json_output:
            print(json.dumps(stats, indent=2))
        else:
            print(format_stats(stats))
        return

    if args.index_only:
        index_data = {
            "directory": args.directory,
            "total_files": len(entries),
            "entries": [
                {
                    "path": e["path"],
                    "title": e["title"],
                    "freshness": e["freshness"],
                    "unique_terms": e["unique_terms"],
                    "sections": e["sections"],
                }
                for e in entries
            ],
        }
        if args.json_output:
            print(json.dumps(index_data, indent=2))
        else:
            print(f"Indexed {len(entries)} files from '{args.directory}'")
            for e in entries:
                marker = {"fresh": "+", "aging": "~", "stale": "-"}.get(e["freshness"], "?")
                print(f"  [{marker}] {e['title']} ({e['unique_terms']} terms) - {e['path']}")
        return

    if args.query:
        results = score_query(args.query, entries, doc_freq)
        if args.json_output:
            output = {
                "query": args.query,
                "total_matches": len(results),
                "results": results[:args.top],
            }
            print(json.dumps(output, indent=2))
        else:
            print(format_search_results(results, args.query, args.top))


if __name__ == "__main__":
    main()