tmux-amq/src/tamq/cli.py

"""Initial tamq command-line surface."""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import signal
import fcntl
from uuid import uuid4
import subprocess
import sys
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path

from . import __version__
from .config import db_path, pid_path, lock_path, config_path, socket_path, state_dir, setting
from .service import Service, ping, request
from .store import Store
from .tmux import TmuxError, TmuxManager
from .broker import BrokerIdentity, InputBroker
from .ptytap import PtyTap
from .control import ControlModeClient
from .diagnostics import configure
from .policy import load_profile


def parse_size(value: str) -> int:
    units = {"B": 1, "KB": 1000, "MB": 1000**2, "GB": 1000**3}
    text = value.upper()
    for unit, multiplier in sorted(units.items(), key=lambda item: -len(item[0])):
        if text.endswith(unit):
            return int(float(text[:-len(unit)]) * multiplier)
    return int(text)


def history_advisory(store: Store) -> str | None:
    size, oldest = store.history_stats()
    limit_text = setting("history_max_size", "100MB")
    limit = parse_size(limit_text)
    if size <= limit:
        return None
    since = time.strftime("%y%m%d", time.localtime(oldest)) if oldest else "unknown"
    return f"Size of history is {size / 1000**2:.0f} MB since {since}, exceeding {limit_text}. Consider using 'tamq purge'."


def ensure_service() -> bool:
    if asyncio.run(ping()):
        return True
    path = lock_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a+") as lock:
        fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
        if asyncio.run(ping()):
            return True
        subprocess.Popen(
            [sys.executable, "-m", "tamq.cli", "serve"],
            stdin=subprocess.DEVNULL,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            start_new_session=True,
            env=os.environ.copy(),
        )
        for _ in range(20):
            if asyncio.run(ping()):
                return True
            time.sleep(0.05)
    return False


def preflight_runtime_paths() -> None:
    """Ensure local state/runtime parents are available before tmux mutation."""
    parents = {db_path().parent, socket_path().parent, pid_path().parent, lock_path().parent}
    for parent in parents:
        parent.mkdir(parents=True, exist_ok=True)
        if not os.access(parent, os.W_OK | os.X_OK):
            raise OSError(f"runtime directory is not writable: {parent}")


def attach_session(session: str) -> int:
    command = ["tmux", "switch-client" if os.environ.get("TMUX") else "attach-session", "-t", session]
    return subprocess.run(command, check=False).returncode


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="tamq",
        description="Tmux Agentic Message Queueing for local repository workers.",
    )
    parser.add_argument("--version", "-V", action="version", version=__version__)
    parser.add_argument("--orwell", action="store_true", help="enable unsafe local diagnostics")
    parser.add_argument("--verbose", action="store_true", help="enable debug diagnostics")
    parser.add_argument("--policy-profile", default=None, help="select configured policy profile")
    subparsers = parser.add_subparsers(dest="command")

    start = subparsers.add_parser("start", help="start or reuse a tamq endpoint")
    start.add_argument("repos", nargs="*", help="gita-registered repository slugs")
    start.add_argument("--command", "--cmd", dest="agent_command", default="codex", help="agent command for new windows (default: codex)")
    start.add_argument("--detach", action="store_true", help="detach after startup")
    start.add_argument("--no-service", action="store_true", help="open tmux windows without service registration or messaging")
    attach = subparsers.add_parser("attach", help="attach to the managed tmux session")
    attach.add_argument("--session", default="tamq", help="tmux session name")

    subparsers.add_parser("serve", help="run the local service in the foreground")
    subparsers.add_parser("stop", help="stop the local service")
    subparsers.add_parser("status", help="show endpoint health")
    subparsers.add_parser("ping", help="check local service liveness")
    history = subparsers.add_parser("history", help="inspect local message history")
    history.add_argument("--repo", dest="target_repo")
    history.add_argument("--state")
    inspect = subparsers.add_parser("inspect", help="inspect one message")
    inspect.add_argument("message_id")
    ack = subparsers.add_parser("ack", help="acknowledge one injected message")
    ack.add_argument("message_id")
    send = subparsers.add_parser("send", help="queue a direct message")
    send.add_argument("address", help="@repo: message")
    send.add_argument("body", nargs="*", help="message body when address is a repo slug")
    send.add_argument("--endpoint-id")
    export = subparsers.add_parser("export", help="export history as JSONL")
    export.add_argument("--output", required=True)
    export.add_argument("--repo", dest="target_repo")
    export.add_argument("--state")
    replay = subparsers.add_parser("replay", help="replay a JSONL message file")
    replay.add_argument("file")
    replay.add_argument("--endpoint-id", help="target a registered tmux-amq endpoint")
    purge = subparsers.add_parser("purge", help="dry-run or remove history")
    purge.add_argument("--before", default=None)
    purge.add_argument("--max-size", default=None)
    purge.add_argument("--yes", action="store_true")
    purge.add_argument("--orwell", action="store_true")
    completion = subparsers.add_parser("completion", help="emit shell completion script")
    completion.add_argument("shell", choices=("bash", "zsh", "fish"))
    subparsers.add_parser("db-version", help="show SQLite schema version")
    subparsers.add_parser("config", help="show effective configuration paths and policy")
    tap = subparsers.add_parser("tap", help="run an agent behind the full-duplex PTY tap")
    tap.add_argument("--repo", required=True, help="source gita repository slug")
    tap.add_argument("--endpoint", required=True, help="tmux-amq endpoint identity")
    tap.add_argument("agent_command", nargs=argparse.REMAINDER, help="agent command after --")
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    configure(args.orwell, args.verbose)
    try:
        profile = load_profile(selected=args.policy_profile or os.environ.get("TAMQ_POLICY_PROFILE", "default"))
    except ValueError as exc:
        print(f"tamq: {exc}", file=sys.stderr)
        return 2
    if args.command is None:
        parser.print_help()
        return 0
    if args.command == "db-version":
        store = Store(db_path()); print(store.db.execute("SELECT value FROM metadata WHERE key='schema_version'").fetchone()[0]); store.close(); return 0
    if args.command == "ping":
        return 0 if asyncio.run(ping()) else 1
    if args.command == "stop":
        try:
            pid = int(pid_path().read_text(encoding="utf-8"))
            os.kill(pid, signal.SIGTERM)
            deadline = time.monotonic() + 2.0
            while time.monotonic() < deadline:
                try:
                    os.kill(pid, 0)
                except ProcessLookupError:
                    break
                time.sleep(0.05)
            print(f"stopped tamq service {pid}")
            return 0
        except ProcessLookupError:
            pid_path().unlink(missing_ok=True)
            print("tamq service is not running", file=sys.stderr)
            return 1
        except (FileNotFoundError, ValueError):
            print("tamq service is not running", file=sys.stderr)
            return 1
    if args.command == "serve":
        try: asyncio.run(Service().run())
        except KeyboardInterrupt: return 0
        return 0
    if args.command == "attach":
        return attach_session(args.session)
    if args.command == "completion":
        print(completion_script(args.shell), end="")
        return 0
    if args.command == "config":
        print(json.dumps({"config": str(config_path()), "state_dir": str(state_dir()), "database": str(db_path()), "socket": str(socket_path()), "pidfile": str(pid_path()), "lockfile": str(lock_path()), "policy_profile": profile.name, "delivery_ack_mode": profile.delivery_ack_mode}, sort_keys=True))
        return 0
    if args.command == "tap":
        command = list(args.agent_command)
        if command and command[0] == "--":
            command = command[1:]
        if not command:
            print("tamq tap requires an agent command after --", file=sys.stderr)
            return 2
        store = Store(db_path())
        try:
            return PtyTap(command, InputBroker(store, BrokerIdentity(args.endpoint, args.repo))).run()
        finally:
            store.close()
    if args.command == "start":
        manager = TmuxManager()
        try:
            preflight_runtime_paths()
            launch_plan = manager.preflight(args.repos, args.agent_command)
        except (TmuxError, OSError) as exc:
            print(f"tamq: {exc}", file=sys.stderr)
            return 2
        if not args.no_service and not ensure_service():
            print("tamq service failed to start; use --no-service to open repos without messaging", file=sys.stderr)
            return 1
        try:
            endpoint = manager.ensure_plan(launch_plan)
        except (TmuxError, OSError) as exc:
            print(f"tamq: {exc}", file=sys.stderr)
            return 2
        delivered = 0
        registered = False
        if not args.no_service:
            try:
                registration = asyncio.run(request({"op": "register", "endpoint_id": endpoint.endpoint_id, "instance_id": endpoint.instance_key, "pid": endpoint.pid, "session": endpoint.session, "repos": endpoint.repos}))
                if not registration.get("ok"):
                    raise RuntimeError(f"endpoint registration failed: {registration.get('error', 'unknown error')}")
                registered = True
                control = ControlModeClient(endpoint.session)
                queue_store = Store(db_path())
                try:
                    control.start()
                    delivered = InputBroker(queue_store, BrokerIdentity(endpoint.instance_key, endpoint.repos[0])).deliver_pending(
                        control, window_for_repo=lambda repo: f"{endpoint.session}:{repo}"
                    )
                finally:
                    control.close()
                    queue_store.close()
            except (OSError, RuntimeError, ValueError, json.JSONDecodeError) as exc:
                manager.rollback(endpoint)
                print(f"tamq: {exc}", file=sys.stderr)
                return 1
        summary = {
            "endpoint_id": endpoint.endpoint_id,
            "instance_id": endpoint.instance_key,
            "session": endpoint.session,
            "repos": endpoint.repos,
            "delivered": delivered,
            "service": not args.no_service,
            "messaging": registered,
        }
        print(json.dumps(summary), flush=True)
        if args.detach:
            return 0
        return attach_session(endpoint.session)
    store = Store(db_path())
    try:
        advisory = history_advisory(store)
        if advisory and args.command in {"start", "serve", "status", "history", "send"}:
            print(advisory, file=sys.stderr)
        if args.command == "send":
            text = " ".join([args.address, *args.body]).strip()
            if not text.startswith("@") or ":" not in text:
                print("tamq send expects @repo: message", file=sys.stderr); return 2
            target, body = text[1:].split(":", 1); body = body.strip()
            if not target or not body: print("tamq send requires a target and body", file=sys.stderr); return 2
            if args.endpoint_id:
                if not asyncio.run(ping()):
                    print("tamq: service is not running", file=sys.stderr); return 1
                response = asyncio.run(request({"op": "send", "endpoint_id": args.endpoint_id, "sender_repo": "local", "target_repo": target, "body": body}))
                if not response.get("ok"):
                    print(f"tamq: {response.get('error', 'send failed')}", file=sys.stderr); return 1
                print(response["message_id"]); return 0
            try:
                print(store.add("local", target, body))
            except ValueError as exc:
                print(f"tamq: {exc}", file=sys.stderr); return 2
            return 0
        if args.command == "history":
            for row in store.list(args.target_repo, args.state): print(json.dumps(dict(row), sort_keys=True))
            return 0
        if args.command == "inspect":
            rows = [row for row in store.list() if row["message_id"] == args.message_id]
            if not rows:
                print("tamq: message not found", file=sys.stderr)
                return 1
            print(json.dumps(dict(rows[0]), sort_keys=True)); return 0
        if args.command == "ack":
            if not store.acknowledge(args.message_id):
                print("tamq: message not found", file=sys.stderr); return 1
            print(args.message_id); return 0
        if args.command == "export":
            store.export(store.list(args.target_repo, args.state), Path(args.output)); return 0
        if args.command == "replay":
            batch_id = f"replay-{uuid4()}"
            count = 0
            if args.endpoint_id and store.endpoint(args.endpoint_id) is None:
                print(f"tamq: endpoint is not registered: {args.endpoint_id}", file=sys.stderr)
                return 1
            try:
                lines = Path(args.file).read_text(encoding="utf-8").splitlines()
                for line in lines:
                    if not line.strip():
                        continue
                    item = json.loads(line)
                    provenance = json.dumps({"batch_id": batch_id, "original_message_id": item.get("message_id")}, sort_keys=True)
                    endpoint = args.endpoint_id or item.get("endpoint_id")
                    store.add(item.get("sender_repo", "replay"), item["target_repo"], item["body"], endpoint=endpoint, provenance=provenance)
                    count += 1
            except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
                print(f"tamq: cannot replay {args.file}: {exc}", file=sys.stderr)
                return 2
            print(json.dumps({"batch_id": batch_id, "count": count})); return 0
        if args.command == "purge":
            before_text = args.before or setting("purge_before", "365d")
            max_size_text = args.max_size or setting("purge_max_size", "100MB")
            before = time.time() - int(before_text[:-1]) * 86400 if before_text.endswith("d") else None
            if not args.yes:
                matching = [row for row in store.list() if before is None or row["created_at"] < before]
                print(f"Dry run: {len(matching)} message(s) match; re-run with --yes to purge.")
                return 0
            count = store.purge(before=before, max_bytes=parse_size(max_size_text))
            print(f"Purged {count} message(s)."); return 0
        if args.command == "status":
            live = asyncio.run(ping())
            endpoints = asyncio.run(request({"op": "endpoints"})) if live else {"endpoints": []}
            print(json.dumps({"version": __version__, "db": str(db_path()), "messages": len(store.list()), "pending": len(store.list(state="pending")), "leases": store.db.execute("SELECT COUNT(*) FROM leases").fetchone()[0], "service": live, "policy_profile": profile.name, "delivery_ack_mode": profile.delivery_ack_mode, "endpoints": endpoints.get("endpoints", [])}, sort_keys=True)); return 0
        print(f"tamq {__version__}: command '{args.command}' is not implemented yet", file=sys.stderr)
        return 2
    finally:
        store.close()


def completion_script(shell: str) -> str:
    commands = "start attach serve stop status ping history inspect ack send export replay purge completion db-version config tap"
    if shell == "bash":
        return f"""_tamq_complete() {{
  local commands=\"{commands}\"
  if [[ ${{COMP_WORDS[1]}} == start || ${{COMP_WORDS[1]}} == send ]]; then
    COMPREPLY=( $(compgen -W \"$(gita ls 2>/dev/null)\" -- \"${{COMP_WORDS[COMP_CWORD]}}\") )
  else
    COMPREPLY=( $(compgen -W \"$commands\" -- \"${{COMP_WORDS[COMP_CWORD]}}\") )
  fi
}}
complete -F _tamq_complete tamq
"""
    if shell == "zsh":
        return f"#compdef tamq\n_arguments '1:command:(({commands}))' '*:repo:->repos'\nif [[ $state == repos ]]; then _values 'gita repo' $(gita ls 2>/dev/null); fi\n"
    return f"complete -c tamq -f -n '__fish_seen_subcommand_from start send' -a '(gita ls 2>/dev/null)'\ncomplete -c tamq -f -a '{commands}'\n"


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