#!/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 _target_label(node=None, container=None):
    if node:
        return f"node netns '{node}' (/tmp/tmux-{node}.sock)"
    if container:
        return f"docker container '{container}'"
    return f"shared socket {SOCKET_PATH}"

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

def _run_tmux(*args, capture=True, node=None, container=None):
    if node:
        cmd = ["/home/super/Projects/NetVM/bin/netvm-exec.sh", node, "--", TMUX_BIN, "-S", f"/tmp/tmux-{node}.sock"] + list(args)
    elif container:
        cmd = ["docker", "exec", "-i", container, "tmux"] + list(args)
    else:
        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)
    node = getattr(args, "node", None)
    container = getattr(args, "container", None)
    
    # Check if session already exists
    rc, out, _ = _run_tmux("has-session", "-t", session, node=node, container=container)
    if rc == 0:
        print(f"Session '{session}' already exists on {_target_label(node, container)}.")
        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, node=node, container=container)
    if rc == 0:
        _ensure_logging(session, node=node, container=container)
        print(f"✔ Created session '{session}' on {_target_label(node, container)} (logging enabled)")
        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)
    node = getattr(args, "node", None)
    container = getattr(args, "container", None)
    
    # Check session
    rc, _, _ = _run_tmux("has-session", "-t", session, node=node, container=container)
    if rc != 0:
        # Auto-create if not exists
        print(f"Notice: Session '{session}' does not exist on {_target_label(node, container)}; creating now...", file=sys.stderr)
        _run_tmux("new-session", "-d", "-s", session, node=node, container=container)
        _ensure_logging(session, node=node, container=container)
    else:
        _ensure_logging(session, node=node, container=container)
        
    t_args = ["send-keys", "-t", session, keys]
    if enter:
        t_args.append("Enter")
        
    rc, out, err = _run_tmux(*t_args, node=node, container=container)
    if rc == 0:
        print(f"✔ Sent keys to '{session}' on {_target_label(node, container)}")
        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
    node = getattr(args, "node", None)
    container = getattr(args, "container", None)
    
    t_args = ["capture-pane", "-p", "-t", session, "-S", f"-{lines}"]
    rc, out, err = _run_tmux(*t_args, node=node, container=container)
    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):
    node = getattr(args, "node", None)
    container = getattr(args, "container", None)
    rc, out, err = _run_tmux("list-sessions", node=node, container=container)
    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 {_target_label(node, container)}.")
        return 0
    else:
        print(f"Error: {err.strip() or out.strip()}", file=sys.stderr)
        return rc

def cmd_kill(args):
    session = args.session
    node = getattr(args, "node", None)
    container = getattr(args, "container", None)
    rc, out, err = _run_tmux("kill-session", "-t", session, node=node, container=container)
    if rc == 0:
        print(f"✔ Killed session '{session}' on {_target_label(node, container)}")
        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
    node = getattr(args, "node", None)
    container = getattr(args, "container", None)
    rc, out, err = _run_tmux("list-sessions", "-F", "#{session_name} #{session_activity} #{session_attached}", node=node, container=container)
    if rc != 0:
        if "no server running" in (err or "").lower() or rc == 1:
            print(f"No active tmux server on {_target_label(node, container)}.")
            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, node=node, container=container)
                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) on {_target_label(node, container)} (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 on {_target_label(node, container)} (all active within {ttl_seconds//3600}h).")
    return 0

def cmd_attach(args):
    session = args.session
    node = getattr(args, "node", None)
    container = getattr(args, "container", None)
    if node:
        cmd = ["/home/super/Projects/NetVM/bin/netvm-exec.sh", node, "--", TMUX_BIN, "-S", f"/tmp/tmux-{node}.sock", "attach", "-t", session]
        os.execv(cmd[0], cmd)
    elif container:
        cmd = ["docker", "exec", "-it", container, "tmux", "attach", "-t", session]
        os.execv("/usr/bin/docker", cmd)
    else:
        cmd = [TMUX_BIN, "-S", SOCKET_PATH, "attach", "-t", session]
        os.execv(TMUX_BIN, cmd)

def main():
    # Common parent parser for hybrid routing
    parent = argparse.ArgumentParser(add_help=False)
    parent.add_argument("--node", default=argparse.SUPPRESS, help="Target NetVM node namespace (e.g. pip, dev, 646, opm, muse, def)")
    parent.add_argument("--container", default=argparse.SUPPRESS, help="Target Docker container name")

    parser = argparse.ArgumentParser(description="Shared Muse tmux socket manager (host, netns nodes, containers)", parents=[parent])
    subparsers = parser.add_subparsers(dest="action")
    
    # list
    p_ls = subparsers.add_parser("list", aliases=["ls"], parents=[parent], help="List sessions on shared socket or node/container")
    p_ls.set_defaults(func=cmd_list)
    
    # new
    p_new = subparsers.add_parser("new", parents=[parent], 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"], parents=[parent], 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"], parents=[parent], 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", parents=[parent], 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", parents=[parent], 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", parents=[parent], 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()