memory.py

"""Save and retrieve learned shell commands.

When a command succeeds it is written to a human-readable markdown file so that
the next identical request skips LLM generation and web research entirely.
Memory is optional and never crashes the app when empty or disabled.
"""

import re
from datetime import date
from pathlib import Path

from config import CONFIG
import llm_client

LEARNED_DIR = Path(__file__).parent / "memory" / "learned_commands"
CACHE_DIR = Path(__file__).parent / "memory" / "research_cache"

FIELD_ORDER = [
    "task",
    "os",
    "command",
    "success_count",
    "last_used",
    "research_required",
    "notes",
]

MAX_SLUG_LENGTH = 60


def ensure_storage():
    """Create the learned-commands and research-cache directories if missing."""
    LEARNED_DIR.mkdir(parents=True, exist_ok=True)
    CACHE_DIR.mkdir(parents=True, exist_ok=True)


def save_command(task, command, os_name, used_research):
    """Save a successful command, incrementing its count if already known."""
    if not CONFIG["memory_enabled"]:
        return
    path = _command_path(task, os_name)
    existing = _read_entry(path)
    entry = _merge_entry(existing, task, command, os_name, used_research)
    if entry is None:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(_serialize(entry))


def find_command(task, os_name):
    """Return a learned command for the task, or None if nothing fits."""
    if not CONFIG["memory_enabled"]:
        return None
    entries = get_all_commands(os_name)
    if not entries:
        return None

    match = _exact_match(task, entries) or _keyword_match(task, entries)
    if match:
        return match["command"]
    return _llm_match(task, entries)


def get_all_commands(os_name):
    """Return all learned command entries for an OS as a list of dicts."""
    directory = LEARNED_DIR / os_name
    if not directory.exists():
        return []
    entries = []
    for path in sorted(directory.glob("*.md")):
        entry = _read_entry(path)
        if entry:
            entries.append(entry)
    return entries


def clear_cache():
    """Delete all research cache files, leaving learned commands untouched."""
    if not CACHE_DIR.exists():
        return 0
    removed = 0
    for path in CACHE_DIR.iterdir():
        if path.is_file():
            path.unlink()
            removed += 1
    return removed


def get_memory_stats():
    """Return totals about learned commands and the research cache."""
    by_os = {}
    total = 0
    oldest = None
    for directory in _os_directories():
        entries = get_all_commands(directory.name)
        by_os[directory.name] = len(entries)
        total += len(entries)
        oldest = _older_of(oldest, entries)
    return {
        "total_learned": total,
        "by_os": by_os,
        "cache_size": _cache_file_count(),
        "oldest_entry": oldest,
    }


def _merge_entry(existing, task, command, os_name, used_research):
    """Build the entry to write, or None to keep an existing command."""
    if existing is None:
        return _new_entry(task, command, os_name, used_research)

    if existing["command"] == command:
        existing["success_count"] += 1
        existing["last_used"] = date.today().isoformat()
        return existing

    return None


def _new_entry(task, command, os_name, used_research):
    """Build a fresh learned-command entry with a success count of one."""
    return {
        "task": task,
        "os": os_name,
        "command": command,
        "success_count": 1,
        "last_used": date.today().isoformat(),
        "research_required": used_research,
        "notes": "",
    }


def _exact_match(task, entries):
    """Return the entry whose task slug equals the request's slug, or None."""
    slug = _slug(task)
    for entry in entries:
        if _slug(entry["task"]) == slug:
            return entry
    return None


def _keyword_match(task, entries):
    """Return the entry with the highest word overlap above the threshold."""
    task_words = _words(task)
    if not task_words:
        return None

    best = None
    best_score = 0.0
    for entry in entries:
        entry_words = _words(entry["task"])
        if not entry_words:
            continue
        score = _jaccard(task_words, entry_words)
        if score > best_score:
            best_score = score
            best = entry

    if best_score >= CONFIG["memory_match_threshold"]:
        return best
    return None


def _jaccard(first, second):
    """Return the Jaccard similarity of two word sets."""
    union = first | second
    if not union:
        return 0.0
    return len(first & second) / len(union)


def _llm_match(task, entries):
    """Ask the LLM if a saved task matches; return its command or None."""
    listing = "\n".join(
        f"{i}. {entry['task']}" for i, entry in enumerate(entries, 1)
    )
    prompt = (
        f"The user wants to: {task}\n\n"
        f"Saved tasks:\n{listing}\n\n"
        "If one saved task achieves the same goal, reply with its number only. "
        "If none match, reply 'none'."
    )
    answer = llm_client.generate(prompt)
    index = _parse_index(answer, len(entries))
    if index is None:
        return None
    return entries[index - 1]["command"]


def _parse_index(answer, count):
    """Return a 1-based index found in the answer within range, or None."""
    if not answer:
        return None
    match = re.search(r"\d+", answer)
    if not match:
        return None
    index = int(match.group())
    return index if 1 <= index <= count else None


def _read_entry(path):
    """Parse a learned command file into a dict with defaults, or None."""
    if not path.exists():
        return None
    try:
        entry = _parse(path.read_text())
    except (OSError, ValueError):
        return None
    return _with_defaults(entry)


def _parse(text):
    """Parse frontmatter 'key: value' lines into a dict."""
    entry = {}
    for line in text.splitlines():
        line = line.strip()
        if not line or line == "---" or ":" not in line:
            continue
        key, value = line.split(":", 1)
        entry[key.strip()] = _coerce(key.strip(), value.strip())
    return entry


def _coerce(key, value):
    """Convert a raw frontmatter value to the right type for its field."""
    if key == "success_count":
        return int(value) if value.isdigit() else 1
    if key == "research_required":
        return value.lower() == "true"
    return value


def _with_defaults(entry):
    """Fill in any missing fields so callers can rely on every key existing."""
    entry.setdefault("task", "")
    entry.setdefault("os", "")
    entry.setdefault("command", "")
    entry.setdefault("success_count", 1)
    entry.setdefault("last_used", date.today().isoformat())
    entry.setdefault("research_required", False)
    entry.setdefault("notes", "")
    return entry


def _serialize(entry):
    """Render an entry dict as a markdown frontmatter block."""
    lines = ["---"]
    for key in FIELD_ORDER:
        lines.append(f"{key}: {_format_value(entry[key])}")
    lines.append("---")
    return "\n".join(lines) + "\n"


def _format_value(value):
    """Render a single field value for the markdown file."""
    if isinstance(value, bool):
        return "true" if value else "false"
    return str(value)


def _command_path(task, os_name):
    """Return the markdown file path for a task on a given OS."""
    return LEARNED_DIR / os_name / f"{_slug(task)}.md"


def _slug(task):
    """Turn a task description into a safe, lowercase filename slug."""
    slug = re.sub(r"[^a-z0-9]+", "_", task.lower()).strip("_")
    return slug[:MAX_SLUG_LENGTH] or "command"


def _words(text):
    """Return the set of meaningful lowercase words in a string."""
    words = re.findall(r"[a-z0-9]+", text.lower())
    return {word for word in words if len(word) > 1}


def _os_directories():
    """Return the existing per-OS directories under learned_commands."""
    if not LEARNED_DIR.exists():
        return []
    return [path for path in LEARNED_DIR.iterdir() if path.is_dir()]


def _older_of(current_oldest, entries):
    """Return the earliest last_used date across the value and entries."""
    oldest = current_oldest
    for entry in entries:
        used = entry["last_used"]
        if used and (oldest is None or used < oldest):
            oldest = used
    return oldest


def _cache_file_count():
    """Return how many files are in the research cache."""
    if not CACHE_DIR.exists():
        return 0
    return sum(1 for path in CACHE_DIR.iterdir() if path.is_file())