#!/usr/bin/env python3
"""Run a command while sampling how the model file is actually read.

Reports, per sample and at the end:
  read_bytes   /proc/<pid>/io          bytes actually fetched from the block device
  file_resident  mincore(model file)   pages of the GGUF sitting in the page cache
  RssAnon/RssFile                     process RSS split (arena+weights vs mapped file)
  Cached          /proc/meminfo       system-wide page cache

The point: a buffered read would show up as file_resident growing; O_DIRECT reads
never touch the cache, so read_bytes grows while file_resident stays flat.

usage: watch_io.py <model.gguf>[,<model2.gguf>,...] <cmd> [args...]
"""
import os
import re
import subprocess
import sys
import time

PAGE = os.sysconf('SC_PAGE_SIZE')
CHR = 1024 ** 3


def meminfo(key):
    for line in open('/proc/meminfo'):
        if line.startswith(key):
            return int(line.split()[1]) * 1024
    return 0


def proc_io(pid):
    out = {}
    try:
        for line in open(f'/proc/{pid}/io'):
            k, _, v = line.partition(':')
            out[k] = int(v)
    except OSError:
        return None
    return out


def proc_status(pid):
    out = {}
    try:
        for line in open(f'/proc/{pid}/status'):
            for k in ('RssAnon', 'RssFile', 'VmHWM'):
                if line.startswith(k):
                    out[k] = int(line.split()[1]) * 1024
    except OSError:
        pass
    return out


def file_resident(paths):
    """resident bytes of the given files via mincore (helper binary)"""
    r = subprocess.run(['/tmp/pagecache_resident'] + paths, capture_output=True, text=True)
    m = re.search(r'TOTAL\s+resident\s+([\d.]+) GiB', r.stdout)
    if m:
        return float(m.group(1)) * CHR
    m = re.search(r'resident\s+([\d.]+) GiB', r.stdout)
    return float(m.group(1)) * CHR if m else -1


def main():
    models = sys.argv[1].split(',')
    cmd = sys.argv[2:]
    subprocess.run(['sync'])
    print(f'[watch] evicting {len(models)} shard(s) from the page cache', flush=True)
    subprocess.run(['/tmp/pagecache_resident', '--evict'] + models)

    base_cached = meminfo('Cached:')
    base_read = None
    t0 = time.time()
    # NOTE: never use pipes here - sdq-chat prints several hundred KB of
    # tokenizer warnings and would block on a full 64 KB pipe.
    out_f = open('/tmp/watch_stdout.log', 'w')
    err_f = open('/tmp/watch_stderr.log', 'w')
    proc = subprocess.Popen(cmd, stdout=out_f, stderr=err_f)
    pid = proc.pid

    peak_res = 0.0
    peak_rssanon = peak_rssfile = 0
    peak_cached = base_cached
    samples = 0
    while proc.poll() is None:
        io_ = proc_io(pid)
        st = proc_status(pid)
        if io_ is None:
            break
        if base_read is None:
            base_read = io_['read_bytes']
        res = file_resident(models)
        cached = meminfo('Cached:')
        peak_res = max(peak_res, res)
        peak_cached = max(peak_cached, cached)
        peak_rssanon = max(peak_rssanon, st.get('RssAnon', 0))
        peak_rssfile = max(peak_rssfile, st.get('RssFile', 0))
        samples += 1
        print(f'[watch] t={time.time()-t0:6.1f}s  read_bytes={(io_["read_bytes"]-base_read)/CHR:7.3f} GiB '
              f' rchar={(io_["rchar"])/CHR:8.3f} GiB  model_file_resident={res/CHR:6.3f} GiB  '
              f'RssAnon={st.get("RssAnon",0)/CHR:6.3f}  RssFile={st.get("RssFile",0)/CHR:6.3f}  '
              f'Cached_delta={(cached-base_cached)/CHR:+7.3f} GiB', flush=True)
        time.sleep(float(os.environ.get("WATCH_INTERVAL", "2.0")))

    proc.wait()
    out_f.close()
    err_f.close()
    print(f'\n[watch] === summary over {samples} samples, {time.time()-t0:.1f}s ===')
    print(f'[watch] peak model_file_resident in page cache : {peak_res/CHR:.3f} GiB')
    print(f'[watch] peak process RssAnon                  : {peak_rssanon/CHR:.3f} GiB')
    print(f'[watch] peak process RssFile                  : {peak_rssfile/CHR:.3f} GiB')
    print(f'[watch] system Cached delta (peak)            : {(peak_cached-base_cached)/CHR:+.3f} GiB')
    print(f'[watch] final model_file_resident             : {file_resident(models)/CHR:.3f} GiB')
    return 0


if __name__ == '__main__':
    sys.exit(main())