"""
gpu_keepalive.py

Tiny AMD GPU keep-alive utility for Windows/WDDM VRAM eviction issues
(e.g. llama.cpp / llama-server on RX 7900 XTX + Adrenalin drivers).

What it does:
  - Opens an OpenCL context on the AMD GPU
  - Allocates a small buffer resident in GPU memory
  - Runs a trivial kernel against it on a timer
  - This keeps the GPU context "active" so the driver's idle
    power-management / residency-trim path doesn't decide to
    evict allocations to system RAM.

This does NOT touch llama.cpp or llama-server directly. It just
keeps the card from ever reaching the idle state that triggers
eviction, so llama-server's own VRAM allocation is less likely
to get trimmed.

Requirements:
    pip install pyopencl

Usage:
    python gpu_keepalive.py                # defaults: 4s interval, 64MB buffer
    python gpu_keepalive.py --interval 3 --size-mb 128
    python gpu_keepalive.py --list         # list available OpenCL devices

Stop with Ctrl+C.
"""

import argparse
import sys
import time

try:
    import pyopencl as cl
    import numpy as np
except ImportError:
    print("Missing dependencies. Install with:\n    pip install pyopencl numpy")
    sys.exit(1)


KERNEL_SRC = """
__kernel void touch(__global float* buf) {
    int i = get_global_id(0);
    buf[i] = buf[i] + 1.0f;
}
"""


def find_amd_device():
    for platform in cl.get_platforms():
        for device in platform.get_devices():
            if device.type & cl.device_type.GPU and "amd" in (platform.name + device.name).lower():
                return platform, device
    # fallback: just grab the first GPU device found
    for platform in cl.get_platforms():
        for device in platform.get_devices():
            if device.type & cl.device_type.GPU:
                return platform, device
    return None, None


def list_devices():
    for platform in cl.get_platforms():
        print(f"Platform: {platform.name}")
        for device in platform.get_devices():
            print(f"  Device: {device.name}  ({device.type})")


def main():
    parser = argparse.ArgumentParser(description="Tiny AMD GPU keep-alive to prevent WDDM VRAM eviction.")
    parser.add_argument("--interval", type=float, default=4.0, help="Seconds between GPU touches (default: 4)")
    parser.add_argument("--size-mb", type=int, default=64, help="Size of the resident buffer in MB (default: 64)")
    parser.add_argument("--list", action="store_true", help="List OpenCL platforms/devices and exit")
    args = parser.parse_args()

    if args.list:
        list_devices()
        return

    platform, device = find_amd_device()
    if device is None:
        print("No OpenCL GPU device found. Is the AMD OpenCL runtime installed (it ships with Adrenalin)?")
        sys.exit(1)

    print(f"Using platform: {platform.name}")
    print(f"Using device:   {device.name}")

    ctx = cl.Context([device])
    queue = cl.CommandQueue(ctx)
    program = cl.Program(ctx, KERNEL_SRC).build()
    kernel = cl.Kernel(program, "touch")

    n_floats = (args.size_mb * 1024 * 1024) // 4
    host_buf = np.zeros(n_floats, dtype=np.float32)
    mf = cl.mem_flags
    gpu_buf = cl.Buffer(ctx, mf.READ_WRITE | mf.COPY_HOST_PTR, hostbuf=host_buf)

    print(f"Resident buffer: {args.size_mb} MB")
    print(f"Touch interval:  {args.interval}s")
    print("Running. Ctrl+C to stop.\n")

    try:
        tick = 0
        while True:
            kernel.set_args(gpu_buf)
            cl.enqueue_nd_range_kernel(queue, kernel, (n_floats,), None)
            queue.finish()
            tick += 1
            print(f"\r[keep-alive] touched GPU buffer #{tick}", end="", flush=True)
            time.sleep(args.interval)
    except KeyboardInterrupt:
        print("\nStopped.")


if __name__ == "__main__":
    main()