"""Tiny multi-turn coding helper with persistent memory.

Companion to the All Things Agentic post 'Memory: Why Your Agent Forgets You'.
Reader chats with `pair` across multiple invocations of this script. Facts the
reader tells `pair` (their stack, project paths, preferences) survive between
runs because they go into a local Chroma vector store under `./.pair-memory/`.

Usage:
    uv sync
    export OPENAI_API_KEY=sk-...
    uv run python pair.py     # session 1: tell it your stack
    uv run python pair.py     # session 2: ask it about your stack

Type :q to exit.

The shell tool is allowlisted to read-only commands. The allowlist is the
safety boundary; do not loosen it without reading the post's footgun section.
"""

from __future__ import annotations

import os
import subprocess
import sys
import uuid
from pathlib import Path

import chromadb
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
from agents import Agent, Runner, SQLiteSession, function_tool


MEMORY_DIR = Path(".pair-memory")
USER_ID = os.environ.get("PAIR_USER", "default")
ALLOWED_COMMANDS = {"ls", "cat", "pytest", "uv"}


# --- memory layer -----------------------------------------------------------

class Memory:
    """Tiny vector-store memory.

    remember(text)        store a fact, embedded by OpenAI text-embedding-3-small.
    recall(query, k=3)    retrieve the k most-relevant prior facts.

    Per-user namespacing is enforced via metadata at write time and a where
    filter at read time. See the post's footgun section for why.
    """

    def __init__(self, persist_dir: Path, user_id: str):
        self.user_id = user_id
        client = chromadb.PersistentClient(path=str(persist_dir))
        embedder = OpenAIEmbeddingFunction(
            api_key=os.environ["OPENAI_API_KEY"],
            model_name="text-embedding-3-small",
        )
        self.collection = client.get_or_create_collection(
            name="pair_memory",
            embedding_function=embedder,
        )

    def remember(self, text: str) -> None:
        self.collection.add(
            ids=[str(uuid.uuid4())],
            documents=[text],
            metadatas=[{"user_id": self.user_id}],
        )

    def recall(self, query: str, k: int = 3) -> list[str]:
        result = self.collection.query(
            query_texts=[query],
            n_results=k,
            where={"user_id": self.user_id},
        )
        return result.get("documents", [[]])[0]


# --- tools (module-global memory so @function_tool callbacks can see it) ----

_memory: Memory | None = None


@function_tool
def save_fact(fact: str) -> str:
    """Store a durable fact about the user (their stack, paths, preferences)."""
    assert _memory is not None
    _memory.remember(fact)
    return "stored"


@function_tool
def lookup(query: str) -> str:
    """Look up prior facts about the user that match the query."""
    assert _memory is not None
    hits = _memory.recall(query, k=3)
    if not hits:
        return "(no matching memory)"
    return "\n".join(f"- {h}" for h in hits)


@function_tool
def run_shell(cmd: str) -> str:
    """Run a read-only shell command. Allowlist: ls, cat, pytest, uv."""
    cmd = cmd.strip()
    head = cmd.split()[0] if cmd else ""
    if head not in ALLOWED_COMMANDS:
        return f"refused: command '{head}' not on the allowlist"
    completed = subprocess.run(
        cmd,
        shell=True,
        capture_output=True,
        text=True,
        timeout=30,
    )
    return (completed.stdout or completed.stderr or "(no output)").strip()


SYSTEM_PROMPT = """You are pair, a small CLI coding helper.

Style:
- One short paragraph or one short bulleted list per reply.
- When the user shares a durable fact (their stack, a project path, a
  preference), call save_fact to remember it.
- Before answering questions about the user's setup, call lookup to fetch
  relevant prior facts.
- Use run_shell for read-only commands the user actually asked for. The
  allowlist is ls, cat, pytest, uv. The tool will refuse anything else.
"""


def main() -> int:
    global _memory
    if not os.environ.get("OPENAI_API_KEY"):
        print("OPENAI_API_KEY is not set", file=sys.stderr)
        return 1

    MEMORY_DIR.mkdir(exist_ok=True)
    _memory = Memory(MEMORY_DIR, user_id=USER_ID)

    agent = Agent(
        name="pair",
        instructions=SYSTEM_PROMPT,
        tools=[run_shell, save_fact, lookup],
        model="gpt-5",
    )
    session = SQLiteSession(":memory:")

    print("pair> (type :q to exit)")
    while True:
        try:
            user_in = input("you> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            return 0
        if user_in in {":q", ":quit", "exit"}:
            return 0
        if not user_in:
            continue
        result = Runner.run_sync(agent, user_in, session=session)
        print(f"pair> {result.final_output}\n")


if __name__ == "__main__":
    raise SystemExit(main())