#!/usr/bin/env python3
"""
Context Overflow Guard Script
Issue #510: [Robustness] Context overflow automation — auto-summarize and commit

Monitors tmux pane context levels and triggers actions at thresholds:
- 60%: Send summarization + commit prompt
- 80%: URGENT force commit, restart fresh with summary
- Logs context levels to tmux-state.json

Usage:
  python3 context-overflow-guard.py              # Run once
  python3 context-overflow-guard.py --daemon     # Run continuously
  python3 context-overflow-guard.py --status     # Show current context levels
"""

import os, sys, json, subprocess, time, re
from datetime import datetime, timezone
from pathlib import Path

# Configuration
LOG_DIR = Path.home() / ".local" / "timmy" / "fleet-health"
STATE_FILE = LOG_DIR / "tmux-state.json"
LOG_FILE = LOG_DIR / "context-overflow.log"

# Thresholds
WARN_THRESHOLD = 60  # % — trigger summarization
URGENT_THRESHOLD = 80  # % — trigger urgent commit

# Skip these sessions
SKIP_SESSIONS = ["Alexander"]

def log(msg):
    """Log message to file and optionally console."""
    timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
    log_entry = "[" + timestamp + "] " + msg
    
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    with open(LOG_FILE, "a") as f:
        f.write(log_entry + "\n")
    
    if "--quiet" not in sys.argv:
        print(log_entry)

def run_tmux(cmd):
    """Run tmux command and return output."""
    try:
        result = subprocess.run(
            "tmux " + cmd,
            shell=True,
            capture_output=True,
            text=True,
            timeout=10
        )
        return result.stdout.strip()
    except Exception as e:
        return ""

def get_sessions():
    """Get all tmux sessions except Alexander."""
    output = run_tmux("list-sessions -F '#{session_name}'")
    if not output:
        return []
    
    sessions = []
    for line in output.split("\n"):
        session = line.strip()
        if session and session not in SKIP_SESSIONS:
            sessions.append(session)
    return sessions

def get_windows(session):
    """Get all windows in a session."""
    output = run_tmux("list-windows -t " + session + " -F '#{window_index}:#{window_name}'")
    if not output:
        return []
    
    windows = []
    for line in output.split("\n"):
        if ":" in line:
            idx, name = line.split(":", 1)
            windows.append({"index": idx, "name": name})
    return windows

def get_panes(session, window_index):
    """Get all panes in a window."""
    target = session + ":" + window_index
    output = run_tmux("list-panes -t " + target + " -F '#{pane_index}'")
    if not output:
        return []
    
    panes = []
    for line in output.split("\n"):
        pane = line.strip()
        if pane:
            panes.append(pane)
    return panes

def capture_pane(session, window_name, pane_index):
    """Capture pane content and extract context info."""
    target = session + ":" + window_name + "." + pane_index
    output = run_tmux("capture-pane -t " + target + " -p 2>&1")
    
    if not output:
        return None
    
    # Look for context bar pattern: ⚕ model | used/total | % | time
    # Example: ⚕ mimo-v2-pro | 45,230/131,072 | 34% | 12m remaining
    context_pattern = r"⚕\s+([^|]+)\|\s*([\d,]+)/([\d,]+)\|\s*(\d+)%\|"
    
    lines = output.split("\n")
    for line in lines:
        match = re.search(context_pattern, line)
        if match:
            model = match.group(1).strip()
            used_str = match.group(2).replace(",", "")
            total_str = match.group(3).replace(",", "")
            percent = int(match.group(4))
            
            try:
                used = int(used_str)
                total = int(total_str)
            except:
                used = 0
                total = 0
            
            return {
                "model": model,
                "used": used,
                "total": total,
                "percent": percent,
                "raw_line": line.strip()
            }
    
    # Alternative pattern: just look for percentage in context-like lines
    percent_pattern = r"(\d+)%"
    for line in lines:
        if "⚕" in line or "remaining" in line.lower() or "context" in line.lower():
            match = re.search(percent_pattern, line)
            if match:
                percent = int(match.group(1))
                return {
                    "model": "unknown",
                    "used": 0,
                    "total": 0,
                    "percent": percent,
                    "raw_line": line.strip()
                }
    
    return None

def send_prompt(session, window_name, pane_index, prompt):
    """Send a prompt to a pane."""
    target = session + ":" + window_name + "." + pane_index
    
    # Escape quotes in prompt
    escaped_prompt = prompt.replace('"', '\\"')
    
    cmd = 'send-keys -t ' + target + ' "/queue ' + escaped_prompt + '" Enter'
    result = run_tmux(cmd)
    
    log("Sent prompt to " + target + ": " + prompt[:50] + "...")
    return result

def restart_pane(session, window_name, pane_index):
    """Restart a pane by sending Ctrl+C twice and restarting hermes."""
    target = session + ":" + window_name + "." + pane_index
    
    # Send Ctrl+C twice to exit
    run_tmux("send-keys -t " + target + " C-c")
    time.sleep(0.5)
    run_tmux("send-keys -t " + target + " C-c")
    time.sleep(1)
    
    # Try to detect profile from process
    pid_cmd = "list-panes -t " + target + " -F '#{pane_pid}'"
    pid = run_tmux(pid_cmd)
    
    if pid:
        # Try to find hermes process with profile
        try:
            ps_result = subprocess.run(
                "ps aux | grep " + pid + " | grep hermes | grep -v grep",
                shell=True,
                capture_output=True,
                text=True,
                timeout=5
            )
            ps_line = ps_result.stdout.strip()
            
            # Look for -p profile flag
            profile_match = re.search(r"-p\s+(\S+)", ps_line)
            if profile_match:
                profile = profile_match.group(1)
                run_tmux("send-keys -t " + target + ' "hermes -p ' + profile + ' chat" Enter')
                log("Restarted pane " + target + " with profile " + profile)
                return
        except:
            pass
    
    # Fallback: just restart with default
    run_tmux("send-keys -t " + target + ' "hermes chat" Enter')
    log("Restarted pane " + target + " with default profile")

def load_state():
    """Load previous state from tmux-state.json."""
    if STATE_FILE.exists():
        try:
            with open(STATE_FILE) as f:
                return json.load(f)
        except:
            pass
    return {"panes": {}, "last_update": None}

def save_state(state):
    """Save state to tmux-state.json."""
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    state["last_update"] = datetime.now(timezone.utc).isoformat()
    
    with open(STATE_FILE, "w") as f:
        json.dump(state, f, indent=2)

def process_pane(session, window_name, pane_index, state):
    """Process a single pane for context overflow."""
    target = session + ":" + window_name + "." + pane_index
    
    # Capture pane
    context_info = capture_pane(session, window_name, pane_index)
    if not context_info:
        return
    
    percent = context_info["percent"]
    
    # Update state
    if "panes" not in state:
        state["panes"] = {}
    
    state["panes"][target] = {
        "context_percent": percent,
        "model": context_info["model"],
        "used": context_info["used"],
        "total": context_info["total"],
        "last_check": datetime.now(timezone.utc).isoformat(),
        "raw_line": context_info["raw_line"]
    }
    
    # Check thresholds
    if percent >= URGENT_THRESHOLD:
        log("URGENT: " + target + " at " + str(percent) + "% — forcing commit and restart")
        
        # Send urgent commit prompt
        urgent_prompt = "URGENT: Context at " + str(percent) + "%. Commit all work NOW, summarize progress, then restart fresh."
        send_prompt(session, window_name, pane_index, urgent_prompt)
        
        # Wait a bit for the prompt to be processed
        time.sleep(2)
        
        # Restart the pane
        restart_pane(session, window_name, pane_index)
        
    elif percent >= WARN_THRESHOLD:
        log("WARN: " + target + " at " + str(percent) + "% — sending summarization prompt")
        
        # Send summarization prompt
        warn_prompt = "Context filling up (" + str(percent) + "%). Summarize current work, commit everything, and prepare for fresh session."
        send_prompt(session, window_name, pane_index, warn_prompt)

def run_once():
    """Run context overflow check once."""
    log("=== Context Overflow Check ===")
    
    state = load_state()
    sessions = get_sessions()
    
    if not sessions:
        log("No tmux sessions found")
        return
    
    total_panes = 0
    warned_panes = 0
    urgent_panes = 0
    
    for session in sessions:
        windows = get_windows(session)
        
        for window in windows:
            window_name = window["name"]
            panes = get_panes(session, window["index"])
            
            for pane_index in panes:
                total_panes += 1
                process_pane(session, window_name, pane_index, state)
                
                target = session + ":" + window_name + "." + pane_index
                if target in state.get("panes", {}):
                    percent = state["panes"][target].get("context_percent", 0)
                    if percent >= URGENT_THRESHOLD:
                        urgent_panes += 1
                    elif percent >= WARN_THRESHOLD:
                        warned_panes += 1
    
    # Save state
    save_state(state)
    
    log("Checked " + str(total_panes) + " panes: " + str(warned_panes) + " warned, " + str(urgent_panes) + " urgent")

def show_status():
    """Show current context levels."""
    state = load_state()
    
    if not state.get("panes"):
        print("No context data available. Run without --status first.")
        return
    
    print("Context Levels (last updated: " + str(state.get("last_update", "unknown")) + ")")
    print("=" * 80)
    
    # Sort by context percentage (highest first)
    panes = sorted(state["panes"].items(), key=lambda x: x[1].get("context_percent", 0), reverse=True)
    
    for target, info in panes:
        percent = info.get("context_percent", 0)
        model = info.get("model", "unknown")
        
        # Color coding
        if percent >= URGENT_THRESHOLD:
            status = "URGENT"
        elif percent >= WARN_THRESHOLD:
            status = "WARN"
        else:
            status = "OK"
        
        print(target.ljust(30) + " " + str(percent).rjust(3) + "% " + status.ljust(7) + " " + model)

def daemon_mode():
    """Run continuously."""
    log("Starting context overflow daemon (check every 60s)")
    
    while True:
        try:
            run_once()
            time.sleep(60)
        except KeyboardInterrupt:
            log("Daemon stopped by user")
            break
        except Exception as e:
            log("Error: " + str(e))
            time.sleep(10)

def main():
    if "--status" in sys.argv:
        show_status()
    elif "--daemon" in sys.argv:
        daemon_mode()
    else:
        run_once()

if __name__ == "__main__":
    main()