cheat.py

#!/usr/bin/env python3
"""
cheat -- Personal command snippet manager. Save, search, run. Zero deps.

Your own cheatsheet: save commands you forget, find them fast, execute them.
Unlike tldr/cheat.sh (community docs), this stores YOUR personal snippets.

Usage:
    py cheat.py add "Start dev server" "py serve.py -p 3000" -t dev,server
    py cheat.py add "Find big files" "py diskuse.py . --bigger 10M"
    py cheat.py find server                       # Search snippets
    py cheat.py find docker                       # Search by keyword
    py cheat.py list                              # Show all snippets
    py cheat.py list -t docker                    # Filter by tag
    py cheat.py run 3                             # Execute snippet #3
    py cheat.py edit 3 --cmd "new command"        # Edit snippet
    py cheat.py rm 3                              # Delete snippet
    py cheat.py tags                              # List all tags
    py cheat.py export > snippets.json            # Export
    py cheat.py import snippets.json              # Import
"""

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

CHEAT_FILE = ".cheatsheet.json"

RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
CYAN = "\033[36m"
RED = "\033[31m"
MAGENTA = "\033[35m"


def color_supported() -> bool:
    if os.environ.get("NO_COLOR"):
        return False
    if sys.platform == "win32":
        return bool(os.environ.get("TERM") or os.environ.get("WT_SESSION"))
    return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()


USE_COLOR = color_supported()


def c(code: str, text: str) -> str:
    return f"{code}{text}{RESET}" if USE_COLOR else text


def get_cheat_path() -> Path:
    """Find cheatsheet file: check CWD, then home directory."""
    local = Path(CHEAT_FILE)
    if local.exists():
        return local
    home = Path.home() / CHEAT_FILE
    if home.exists():
        return home
    return home  # default to home


def load(path: Path = None) -> list[dict]:
    if path is None:
        path = get_cheat_path()
    if not path.exists():
        return []
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, Exception):
        return []


def save(snippets: list[dict], path: Path = None):
    if path is None:
        path = get_cheat_path()
    path.write_text(json.dumps(snippets, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


def next_id(snippets: list[dict]) -> int:
    if not snippets:
        return 1
    return max(s.get("id", 0) for s in snippets) + 1


def format_snippet(s: dict, show_cmd: bool = True):
    """Display a single snippet."""
    sid = s.get("id", 0)
    desc = s.get("desc", "")
    cmd = s.get("cmd", "")
    tags = s.get("tags", [])
    tag_str = " " + " ".join(c(MAGENTA, f"[{t}]") for t in tags) if tags else ""

    print(f"  {c(YELLOW, f'#{sid:<4}')} {c(BOLD, desc)}{tag_str}")
    if show_cmd:
        print(f"         {c(CYAN, cmd)}")


def search_snippets(snippets: list[dict], query: str) -> list[dict]:
    """Search snippets by keyword across desc, cmd, and tags."""
    query_lower = query.lower()
    terms = query_lower.split()
    scored = []
    for s in snippets:
        searchable = f"{s.get('desc', '')} {s.get('cmd', '')} {' '.join(s.get('tags', []))}".lower()
        score = sum(2 if t in s.get("desc", "").lower() else 1 for t in terms if t in searchable)
        if score > 0:
            scored.append((score, s))
    scored.sort(key=lambda x: -x[0])
    return [s for _, s in scored]


# --- Commands ---

def cmd_add(args):
    snippets = load()
    snippet = {
        "id": next_id(snippets),
        "desc": args.desc,
        "cmd": args.cmd,
        "tags": [t.strip() for t in args.tags.split(",")] if args.tags else [],
    }
    snippets.append(snippet)
    save(snippets)
    print(f"  {c(GREEN, 'Added')} snippet #{snippet['id']}")
    format_snippet(snippet)


def cmd_find(args):
    snippets = load()
    if not snippets:
        print(c(DIM, "  No snippets. Add one: py cheat.py add \"desc\" \"command\""))
        return

    results = search_snippets(snippets, args.query)
    if not results:
        print(f"  No snippets matching '{args.query}'")
        return

    print(f"\n  {c(BOLD, f'Search: \"{args.query}\"')} ({len(results)} results)\n")
    for s in results:
        format_snippet(s)
    print()


def cmd_list(args):
    snippets = load()
    if not snippets:
        print(c(DIM, "  No snippets saved yet."))
        print(c(DIM, "  Add one: py cheat.py add \"description\" \"command\" -t tag"))
        return

    # Filter by tag
    if args.tag:
        snippets = [s for s in snippets if args.tag in s.get("tags", [])]
        if not snippets:
            print(f"  No snippets with tag '{args.tag}'")
            return

    print(f"\n  {c(BOLD, f'Snippets')} ({len(snippets)})\n")
    for s in snippets:
        format_snippet(s)
    print()


def cmd_run(args):
    snippets = load()
    target = None

    # Find by ID
    for s in snippets:
        if s.get("id") == args.id:
            target = s
            break

    if not target:
        print(c(RED, f"  Snippet #{args.id} not found."))
        return

    cmd = target["cmd"]
    print(f"  {c(DIM, 'Running:')} {c(CYAN, cmd)}")
    print(c(DIM, "  " + "-" * 50))
    result = subprocess.run(cmd, shell=True)
    print(c(DIM, "  " + "-" * 50))
    if result.returncode == 0:
        print(c(GREEN, f"  OK (exit 0)"))
    else:
        print(c(RED, f"  Exit code: {result.returncode}"))


def cmd_edit(args):
    snippets = load()
    for s in snippets:
        if s.get("id") == args.id:
            if args.desc:
                s["desc"] = args.desc
            if args.cmd:
                s["cmd"] = args.cmd
            if args.tags is not None:
                s["tags"] = [t.strip() for t in args.tags.split(",")] if args.tags else []
            save(snippets)
            print(c(GREEN, f"  Updated snippet #{args.id}"))
            format_snippet(s)
            return
    print(c(RED, f"  Snippet #{args.id} not found."))


def cmd_rm(args):
    snippets = load()
    new = [s for s in snippets if s.get("id") != args.id]
    if len(new) == len(snippets):
        print(c(RED, f"  Snippet #{args.id} not found."))
        return
    save(new)
    print(c(GREEN, f"  Removed snippet #{args.id}"))


def cmd_tags(args):
    snippets = load()
    tag_counts: dict[str, int] = {}
    for s in snippets:
        for t in s.get("tags", []):
            tag_counts[t] = tag_counts.get(t, 0) + 1

    if not tag_counts:
        print(c(DIM, "  No tags yet."))
        return

    print(f"\n  {c(BOLD, 'Tags')}\n")
    for tag, count in sorted(tag_counts.items(), key=lambda x: -x[1]):
        print(f"  {c(MAGENTA, tag):<20} {count} snippet(s)")
    print()


def cmd_export(args):
    snippets = load()
    print(json.dumps(snippets, indent=2, ensure_ascii=False))


def cmd_import(args):
    path = Path(args.file)
    if not path.exists():
        print(c(RED, f"  File not found: {args.file}"))
        return
    try:
        incoming = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as e:
        print(c(RED, f"  Invalid JSON: {e}"))
        return

    snippets = load()
    existing_ids = {s.get("id") for s in snippets}

    added = 0
    for s in incoming:
        # Re-assign IDs to avoid conflicts
        s["id"] = next_id(snippets)
        snippets.append(s)
        added += 1

    save(snippets)
    print(c(GREEN, f"  Imported {added} snippet(s). Total: {len(snippets)}"))


def main():
    parser = argparse.ArgumentParser(
        description="cheat -- personal command snippet manager",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    sub = parser.add_subparsers(dest="command")

    p = sub.add_parser("add", help="Add a snippet")
    p.add_argument("desc", help="Description")
    p.add_argument("cmd", help="Command to save")
    p.add_argument("-t", "--tags", help="Comma-separated tags")

    p = sub.add_parser("find", help="Search snippets")
    p.add_argument("query", help="Search term")

    p = sub.add_parser("list", aliases=["ls"], help="List all snippets")
    p.add_argument("-t", "--tag", help="Filter by tag")

    p = sub.add_parser("run", help="Execute a snippet by ID")
    p.add_argument("id", type=int, help="Snippet ID")

    p = sub.add_parser("edit", help="Edit a snippet")
    p.add_argument("id", type=int, help="Snippet ID")
    p.add_argument("--desc", help="New description")
    p.add_argument("--cmd", help="New command")
    p.add_argument("--tags", help="New tags (comma-separated, empty to clear)")

    p = sub.add_parser("rm", help="Delete a snippet")
    p.add_argument("id", type=int, help="Snippet ID")

    sub.add_parser("tags", help="List all tags with counts")
    sub.add_parser("export", help="Export snippets as JSON")

    p = sub.add_parser("import", help="Import snippets from JSON file")
    p.add_argument("file", help="JSON file to import")

    args = parser.parse_args()
    cmds = {
        "add": cmd_add, "find": cmd_find, "list": cmd_list, "ls": cmd_list,
        "run": cmd_run, "edit": cmd_edit, "rm": cmd_rm, "tags": cmd_tags,
        "export": cmd_export, "import": cmd_import,
    }

    if args.command in cmds:
        cmds[args.command](args)
    else:
        # Default: list
        cmd_list(argparse.Namespace(tag=None))


if __name__ == "__main__":
    main()