#!/usr/bin/env python3
"""
muse-tmux.py: Shared tmux socket manager for Muse agents and operators.
Socket location: /tmp/tmux-muse.sock (shared across fleet agents & super).
"""
import sys
import os
import re
import subprocess
import argparse

SOCKET_PATH = "/tmp/tmux-muse.sock"
LOG_DIR = "/home/super/Projects/NetVM/logs/tmux"
TMUX_BIN = "/home/super/.local/bin/tmux"
if not os.path.exists(TMUX_BIN):
    TMUX_BIN = "tmux"

def _ensure_logging(session):
    try:
        os.makedirs(LOG_DIR, exist_ok=True)
        log_file = os.path.join(LOG_DIR, f"{session}.log")
        _run_tmux("pipe-pane", "-o", "-t", session, f"cat >> {log_file}")
    except Exception:
        pass

def _run_tmux(*args, capture=True):
    cmd = [TMUX_BIN, "-S", SOCKET_PATH] + list(args)
    res = subprocess.run(cmd, capture_output=capture, text=True)
    return res.returncode, res.stdout, res.stderr

def cmd_new(args):
    session = args.session
    window = getattr(args, "window", None)
    command = getattr(args, "command", None)
    
    # Check if session already exists
    rc, out, _ = _run_tmux("has-session", "-t", session)
    if rc == 0:
        print(f"Session '{session}' already exists on shared socket {SOCKET_PATH}.")
        return 0
        
    t_args = ["new-session", "-d", "-s", session]
    if window:
        t_args.extend(["-n", window])
    if command:
        t_args.append(command)
        
    rc, out, err = _run_tmux(*t_args)
    if rc == 0:
        _ensure_logging(session)
        print(f"✔ Created session '{session}' on {SOCKET_PATH} (logging to logs/tmux/{session}.log)")
        return 0
    else:
        print(f"Failed to create session: {err.strip() or out.strip()}", file=sys.stderr)
        return rc

def cmd_send(args):
    session = args.session
    keys = args.keys
    enter = not getattr(args, "no_enter", False)
    
    # Check session
    rc, _, _ = _run_tmux("has-session", "-t", session)
    if rc != 0:
        # Auto-create if not exists
        print(f"Notice: Session '{session}' does not exist; creating now...", file=sys.stderr)
        _run_tmux("new-session", "-d", "-s", session)
        _ensure_logging(session)
    else:
        _ensure_logging(session)
        
    t_args = ["send-keys", "-t", session, keys]
    if enter:
        t_args.append("Enter")
        
    rc, out, err = _run_tmux(*t_args)
    if rc == 0:
        print(f"✔ Sent keys to '{session}'")
        return 0
    else:
        print(f"Failed to send keys: {err.strip() or out.strip()}", file=sys.stderr)
        return rc

def cmd_capture(args):
    session = args.session
    lines = getattr(args, "lines", 30) or 30
    
    t_args = ["capture-pane", "-p", "-t", session, "-S", f"-{lines}"]
    rc, out, err = _run_tmux(*t_args)
    if rc == 0:
        # Strip trailing blank lines
        clean = out.rstrip()
        print(clean)
        return 0
    else:
        print(f"Failed to capture pane: {err.strip() or out.strip()}", file=sys.stderr)
        return rc

def cmd_list(args):
    rc, out, err = _run_tmux("list-sessions")
    if rc == 0:
        print(out.strip() or "No active sessions.")
        return 0
    elif "no server running" in (err or "").lower() or rc == 1:
        print(f"No active tmux server on {SOCKET_PATH}.")
        return 0
    else:
        print(f"Error: {err.strip() or out.strip()}", file=sys.stderr)
        return rc

def cmd_kill(args):
    session = args.session
    rc, out, err = _run_tmux("kill-session", "-t", session)
    if rc == 0:
        print(f"✔ Killed session '{session}'")
        return 0
    else:
        print(f"Error: {err.strip() or out.strip()}", file=sys.stderr)
        return rc

def cmd_prune(args):
    ttl_seconds = getattr(args, "ttl", 7200) or 7200
    rc, out, err = _run_tmux("list-sessions", "-F", "#{session_name} #{session_activity} #{session_attached}")
    if rc != 0:
        if "no server running" in (err or "").lower() or rc == 1:
            print(f"No active tmux server on {SOCKET_PATH}.")
            return 0
        print(f"Error listing sessions: {err.strip() or out.strip()}", file=sys.stderr)
        return rc
        
    now = int(time.time())
    reaped = []
    kept = []
    
    for line in out.strip().splitlines():
        parts = line.split()
        if len(parts) >= 3:
            s_name, s_activity, s_attached = parts[0], int(parts[1]), int(parts[2])
            idle_time = now - s_activity
            # Only prune if unattached and idle > TTL
            if s_attached == 0 and idle_time > ttl_seconds:
                _run_tmux("kill-session", "-t", s_name)
                reaped.append((s_name, idle_time))
            else:
                kept.append((s_name, idle_time, s_attached))
                
    if reaped:
        print(f"✔ Reaped {len(reaped)} stale unattached session(s) (idle > {ttl_seconds//3600}h):")
        for s_name, idle in reaped:
            print(f"  • {s_name} (idle: {idle//60}m)")
    else:
        print(f"No stale unattached sessions to reap (all active within {ttl_seconds//3600}h).")
    return 0

def cmd_attach(args):
    session = args.session
    cmd = [TMUX_BIN, "-S", SOCKET_PATH, "attach", "-t", session]
    # Replace current process for interactive attach
    os.execv(TMUX_BIN, cmd)

def main():
    parser = argparse.ArgumentParser(description="Shared Muse tmux socket manager")
    subparsers = parser.add_subparsers(dest="action")
    
    # list
    p_ls = subparsers.add_parser("list", aliases=["ls"], help="List sessions on shared socket")
    p_ls.set_defaults(func=cmd_list)
    
    # new
    p_new = subparsers.add_parser("new", help="Create new background session")
    p_new.add_argument("session", help="Session name")
    p_new.add_argument("--window", "-w", help="Initial window name")
    p_new.add_argument("--command", "-c", help="Command to run in session")
    p_new.set_defaults(func=cmd_new)
    
    # send / send-keys
    p_send = subparsers.add_parser("send", aliases=["send-keys"], help="Send keys to a session")
    p_send.add_argument("session", help="Target session name")
    p_send.add_argument("keys", help="Keys / command string to send")
    p_send.add_argument("--no-enter", action="store_true", help="Do not send Enter key after keys")
    p_send.set_defaults(func=cmd_send)
    
    # capture
    p_cap = subparsers.add_parser("capture", aliases=["cap", "tail"], help="Capture pane output")
    p_cap.add_argument("session", help="Target session name")
    p_cap.add_argument("--lines", "-n", type=int, default=30, help="Number of scrollback lines to capture (default: 30)")
    p_cap.set_defaults(func=cmd_capture)
    
    # kill
    p_kill = subparsers.add_parser("kill", help="Kill a session")
    p_kill.add_argument("session", help="Session name")
    p_kill.set_defaults(func=cmd_kill)
    
    # prune
    p_prune = subparsers.add_parser("prune", help="Reap stale unattached sessions inactive for >TTL (default: 7200s / 2h)")
    p_prune.add_argument("--ttl", type=int, default=7200, help="Inactivity threshold in seconds (default: 7200)")
    p_prune.set_defaults(func=cmd_prune)
    
    # attach
    p_att = subparsers.add_parser("attach", help="Attach to a session interactively")
    p_att.add_argument("session", help="Session name")
    p_att.set_defaults(func=cmd_attach)
    
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(0)
        
    args = parser.parse_args()
    if not hasattr(args, "func"):
        parser.print_help()
        sys.exit(1)
        
    sys.exit(args.func(args))

if __name__ == "__main__":
    main()