#!/usr/bin/env python3
"""
Drop-in replacement for w(1).
procps-ng 4.x w never reads utmp, giving wrong TTY, IDLE, and WHAT.
Reads utmp via who(1), idle via tty atime, foreground pgrp via /proc/pid/stat tpgid.
"""
import os
import re
import subprocess
import time
from datetime import datetime

CLK_TCK = os.sysconf('SC_CLK_TCK')


def fmt_idle(secs):
    if secs is None:
        return '?'
    s = max(0, int(secs))
    if s < 60:
        return f'{s}s'
    m, s2 = divmod(s, 60)
    if m < 60:
        return f'{m}:{s2:02d}'
    h, m2 = divmod(m, 60)
    return f'{h}:{m2:02d}h' if h < 24 else f'{h // 24}day'


def fmt_cpu(ticks):
    secs = ticks / CLK_TCK
    if secs < 60:
        return f'{secs:.1f}s'
    m, s = divmod(secs, 60)
    if m < 60:
        return f'{int(m)}:{s:04.1f}'
    h, m2 = divmod(int(m), 60)
    return f'{h}:{m2:02d}:{int(s):02d}'


def tty_name_to_nr(name):
    """
    Convert a tty name to the integer used in /proc/pid/stat field 7.
    Encoding: bits 0-7 = minor[0:7], bits 8-19 = major, bits 20-31 = minor[8:19].
    """
    if name.startswith('pts/') and name[4:].isdigit():
        n = int(name[4:])
        return (136 << 8) | (n & 0xff) | ((n >> 8) << 20)
    if name.startswith('tty') and name[3:].isdigit():
        return (4 << 8) | int(name[3:])
    return None


# Single /proc scan — build everything we need.
tty_tpgid   = {}   # tty_nr -> foreground pgrp (tpgid from /proc/pid/stat)
tty_session = {}   # tty_nr -> session leader pid  (only when session == pid)
session_cpu = {}   # session leader pid -> cumulative cpu ticks
pgrp_cpu    = {}   # pgrp -> cumulative cpu ticks
pgrp_cmd    = {}   # pgrp -> first non-empty cmdline found for that pgrp

for pid_s in os.listdir('/proc'):
    if not pid_s.isdigit():
        continue
    pid = int(pid_s)
    try:
        raw  = open(f'/proc/{pid}/stat').read()
        ce   = raw.rindex(')')
        flds = raw[ce + 2:].split()
        pgrp    = int(flds[2])
        session = int(flds[3])
        tty_nr  = int(flds[4])
        tpgid   = int(flds[5])
        cpu     = int(flds[11]) + int(flds[12])   # utime + stime

        session_cpu[session] = session_cpu.get(session, 0) + cpu
        pgrp_cpu[pgrp]       = pgrp_cpu.get(pgrp, 0) + cpu

        if tty_nr > 0:
            tty_tpgid[tty_nr] = tpgid
            if session == pid:   # session leaders only
                tty_session[tty_nr] = pid

        if pgrp not in pgrp_cmd:
            try:
                cmd = open(f'/proc/{pid}/cmdline').read().replace('\0', ' ').strip()
                if cmd:
                    pgrp_cmd[pgrp] = cmd
            except OSError:
                pass
    except OSError:
        pass


# Parse who -u: USER TTY YYYY-MM-DD HH:MM IDLE PID [(FROM)]
who_out = subprocess.run(['who', '-u'], capture_output=True, text=True).stdout
pat = re.compile(
    r'^(\S+)\s+(\S+)\s+(\d{4}-\d{2}-\d{2})\s+(\d{2}:\d{2})\s+\S+\s+\d+'
    r'(?:\s+\(([^)]*)\))?'
)
users = []
for line in who_out.splitlines():
    m = pat.match(line)
    if m:
        user, tty, date, hhmm, from_ = m.groups()
        login_ts = datetime.strptime(f'{date} {hhmm}', '%Y-%m-%d %H:%M').timestamp()
        users.append((user, tty, login_ts, from_ or ''))


now = time.time()
print(subprocess.run(['uptime'], capture_output=True, text=True).stdout.strip())
print(f'{"USER":<10} {"TTY":<8} {"FROM":<16} {"LOGIN@":<7} {"IDLE":<6} {"JCPU":<7} {"PCPU":<7} WHAT')

for user, tty, login_ts, from_ in users:
    tty_nr = tty_name_to_nr(tty)

    idle = None
    try:
        idle = now - os.stat(f'/dev/{tty}').st_atime
    except OSError:
        pass

    sess  = tty_session.get(tty_nr, -1) if tty_nr else -1
    jcpu  = session_cpu.get(sess, 0) if sess > 0 else 0

    fpgrp = tty_tpgid.get(tty_nr, -1) if tty_nr else -1
    pcpu  = pgrp_cpu.get(fpgrp, 0) if fpgrp > 0 else 0
    what  = pgrp_cmd.get(fpgrp, '-') if fpgrp > 0 else '-'

    login_str = datetime.fromtimestamp(login_ts).strftime('%H:%M')
    print(
        f'{user:<10} {tty:<8} {from_:<16} {login_str:<7}'
        f' {fmt_idle(idle):<6} {fmt_cpu(jcpu):<7} {fmt_cpu(pcpu):<7} {what[:48]}'
    )
