#!/usr/bin/env python3
"""
watch_remote_top.py

Prompts for a remote host/IP, SSH username, and password (interactively,
password not echoed), then every 10 minutes:
  - SSHes in and runs `top -bn1 -o %MEM` (one-shot, non-interactive, sorted by RAM)
  - Saves the full raw output to a local log file
  - Parses out CPU usage, RAM usage, and the top RAM-consuming process into a
    local CSV summary

Everything is saved on THIS machine, not the remote host.

Requires: paramiko
    pip install paramiko

Usage:
    python3 watch_remote_top.py
    python3 watch_remote_top.py --once            # single run, no loop
    python3 watch_remote_top.py --interval 5       # poll every 5 minutes instead of 10
    python3 watch_remote_top.py --log-dir ./logs   # where to save logs (default: ./TopLogs)
"""

import argparse
import csv
import getpass
import re
import sys
import time
from datetime import datetime
from pathlib import Path

try:
    import paramiko
except ImportError:
    sys.exit(
        "The 'paramiko' package is required but not installed.\n"
        "Install it with:  pip install paramiko"
    )

# Plain `-bn1` (no -o): BusyBox's top doesn't support -o for sorting, so we
# capture unsorted output and sort by memory ourselves in parse_top_output().
REMOTE_COMMAND = "top -bn1"


def get_connection_details():
    host = input("Remote host/IP: ").strip()
    username = input("SSH username: ").strip()
    password = getpass.getpass("SSH password: ")
    return host, username, password


def run_remote_top(host, username, password, timeout=10):
    """SSH in, run the top command once, return the raw stdout as a string."""
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        client.connect(
            hostname=host,
            username=username,
            password=password,
            timeout=timeout,
            banner_timeout=timeout,
            auth_timeout=timeout,
        )
        stdin, stdout, stderr = client.exec_command(REMOTE_COMMAND, timeout=timeout)
        output = stdout.read().decode(errors="replace")
        err = stderr.read().decode(errors="replace")
        if not output and err:
            raise RuntimeError(f"Remote command produced no output. stderr: {err.strip()}")
        return output
    finally:
        client.close()


def _to_mib(value, unit):
    """Convert a K/M/G-suffixed value (as used by BusyBox's `Mem:` line) to MiB."""
    value = float(value)
    unit = (unit or "").upper()
    if unit.startswith("K"):
        return value / 1024
    if unit.startswith("G"):
        return value * 1024
    return value  # already M / MiB, or no unit


def parse_top_output(raw_output):
    """
    Pull CPU used %, mem used/total/free (MiB), mem used %, and the top
    RAM-consuming process out of raw `top -bn1` output. Handles both the
    standard Linux (procps) format and BusyBox's leaner top format, since
    the remote command has no reliable way to sort by memory itself
    (BusyBox top doesn't support -o), so the top process is found by
    scanning and sorting the process rows here instead.
    """
    lines = raw_output.splitlines()

    cpu_used = cpu_idle = None
    mem_total = mem_used = mem_free = mem_used_pct = None
    top_process = None

    is_procps = any("%Cpu(s)" in line for line in lines)
    is_busybox = any(re.match(r"^CPU:", line.strip()) for line in lines)

    if is_procps:
        for line in lines:
            if "%Cpu(s)" in line:
                m = re.search(r"([\d.]+)\s*id", line)
                if m:
                    cpu_idle = float(m.group(1))
                    cpu_used = round(100 - cpu_idle, 1)

            if re.search(r"(MiB|KiB|GiB)\s+Mem", line):
                m_total = re.search(r"([\d.]+)\s+total", line)
                m_free = re.search(r"([\d.]+)\s+free", line)
                m_used = re.search(r"([\d.]+)\s+used", line)
                if m_total:
                    mem_total = float(m_total.group(1))
                if m_free:
                    mem_free = float(m_free.group(1))
                if m_used:
                    mem_used = float(m_used.group(1))
        if mem_total and mem_used:
            mem_used_pct = round((mem_used / mem_total) * 100, 1)

        # PID USER PR NI VIRT RES SHR S %CPU %MEM TIME+ COMMAND
        header_idx, mem_col_idx = None, None
        for i, line in enumerate(lines):
            if re.match(r"^\s*PID\s+USER", line):
                header_cols = line.split()
                if "%MEM" in header_cols:
                    header_idx = i
                    mem_col_idx = header_cols.index("%MEM")
                break

        if header_idx is not None and mem_col_idx is not None:
            best_pct = -1.0
            for line in lines[header_idx + 1:]:
                cols = line.split()
                if len(cols) <= mem_col_idx:
                    continue
                try:
                    pct = float(cols[mem_col_idx])
                except ValueError:
                    continue
                if pct > best_pct:
                    best_pct = pct
                    top_process = f"{cols[-1]} ({pct}% MEM)"

    elif is_busybox:
        # e.g. "Mem: 51380K used, 178308K free, 0K shrd, 3844K buff, 21520K cached"
        for line in lines:
            m = re.search(
                r"Mem:\s*([\d.]+)\s*([KMG])?\s*used,\s*([\d.]+)\s*([KMG])?\s*free",
                line,
                re.IGNORECASE,
            )
            if m:
                mem_used = round(_to_mib(m.group(1), m.group(2)), 1)
                mem_free = round(_to_mib(m.group(3), m.group(4)), 1)
                mem_total = round(mem_used + mem_free, 1)
                if mem_total:
                    mem_used_pct = round((mem_used / mem_total) * 100, 1)

            # e.g. "CPU:   2.3% usr   1.1% sys   0.0% nic  96.4% idle ..."
            m_idle = re.search(r"([\d.]+)%?\s*idle", line, re.IGNORECASE)
            if m_idle:
                cpu_idle = float(m_idle.group(1))
                cpu_used = round(100 - cpu_idle, 1)

        # e.g. "  PID  PPID USER     STAT   VSZ %VSZ %CPU COMMAND"
        header_idx, mem_col_idx = None, None
        for i, line in enumerate(lines):
            cols = line.split()
            if cols[:1] == ["PID"] and "USER" in cols:
                header_idx = i
                if "%VSZ" in cols:
                    mem_col_idx = cols.index("%VSZ")
                elif "%MEM" in cols:
                    mem_col_idx = cols.index("%MEM")
                break

        if header_idx is not None and mem_col_idx is not None:
            best_pct = -1.0
            for line in lines[header_idx + 1:]:
                cols = line.split()
                if len(cols) <= mem_col_idx:
                    continue
                try:
                    pct = float(cols[mem_col_idx].rstrip("%"))
                except ValueError:
                    continue
                if pct > best_pct:
                    best_pct = pct
                    top_process = f"{cols[-1]} ({pct}% VSZ)"

    return {
        "cpu_used": cpu_used,
        "cpu_idle": cpu_idle,
        "mem_total": mem_total,
        "mem_used": mem_used,
        "mem_free": mem_free,
        "mem_used_pct": mem_used_pct,
        "top_process": top_process,
    }


def ensure_csv_header(csv_path):
    if not csv_path.exists():
        with open(csv_path, "w", newline="", encoding="utf-8") as f:
            writer = csv.writer(f)
            writer.writerow(
                [
                    "Timestamp",
                    "CPU_Used_Pct",
                    "CPU_Idle_Pct",
                    "Mem_Total_MiB",
                    "Mem_Used_MiB",
                    "Mem_Free_MiB",
                    "Mem_Used_Pct",
                    "Top_RAM_Process",
                ]
            )


def collect_snapshot(host, username, password, full_log_path, summary_csv_path):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] Connecting to {host}...")

    try:
        raw_output = run_remote_top(host, username, password)
    except Exception as exc:
        print(f"[{timestamp}] SSH/command failed: {exc}", file=sys.stderr)
        with open(summary_csv_path, "a", newline="", encoding="utf-8") as f:
            csv.writer(f).writerow([timestamp, "ERROR", "", "", "", "", "", f"SSH failed: {exc}"])
        return

    # Save full raw snapshot
    with open(full_log_path, "a", encoding="utf-8") as f:
        f.write(f"===== {timestamp} =====\n")
        f.write(raw_output)
        f.write("\n\n")

    parsed = parse_top_output(raw_output)

    with open(summary_csv_path, "a", newline="", encoding="utf-8") as f:
        csv.writer(f).writerow(
            [
                timestamp,
                parsed["cpu_used"],
                parsed["cpu_idle"],
                parsed["mem_total"],
                parsed["mem_used"],
                parsed["mem_free"],
                parsed["mem_used_pct"],
                parsed["top_process"],
            ]
        )

    print(
        f"[{timestamp}] CPU used: {parsed['cpu_used']}% | "
        f"Mem used: {parsed['mem_used_pct']}% "
        f"({parsed['mem_used']}/{parsed['mem_total']} MiB) | "
        f"Top RAM proc: {parsed['top_process']}"
    )


def main():
    parser = argparse.ArgumentParser(description="Poll a remote Linux host's `top` output over SSH and log it locally.")
    parser.add_argument("--once", action="store_true", help="Run a single collection and exit instead of looping.")
    parser.add_argument("--interval", type=int, default=10, help="Polling interval in minutes (default: 10).")
    parser.add_argument("--log-dir", type=str, default="TopLogs", help="Local folder to save logs in (default: ./TopLogs).")
    args = parser.parse_args()

    host, username, password = get_connection_details()

    log_dir = Path(args.log_dir)
    log_dir.mkdir(parents=True, exist_ok=True)
    full_log_path = log_dir / "top_full_snapshots.log"
    summary_csv_path = log_dir / "cpu_ram_summary.csv"
    ensure_csv_header(summary_csv_path)

    print(f"\nFull snapshots -> {full_log_path}")
    print(f"Summary CSV    -> {summary_csv_path}\n")

    if args.once:
        collect_snapshot(host, username, password, full_log_path, summary_csv_path)
        return

    print(f"Polling {host} every {args.interval} minute(s). Press Ctrl+C to stop.\n")
    try:
        while True:
            collect_snapshot(host, username, password, full_log_path, summary_csv_path)
            time.sleep(args.interval * 60)
    except KeyboardInterrupt:
        print("\nStopped.")


if __name__ == "__main__":
    main()
