#!/usr/bin/env python3

import grp
import os
import pathlib
import re
import socket
import subprocess
import sys
import time

try:
    import psutil
except ImportError:
    print("sudo apt install python3-psutil")
    sys.exit(1)

try:
    import requests
except ImportError:
    print("sudo apt install python3-requests")
    sys.exit(1)


DOCKER_DEVICE_ARGS_BY_GPU = {
    "amd": [
        "--device", "/dev/dri",
        "--group-add", str(grp.getgrnam("video").gr_gid),
        "--group-add", str(grp.getgrnam("render").gr_gid),
    ],
    "intel": [
        "--device", "/dev/dri",
        "--group-add", str(grp.getgrnam("video").gr_gid),
        "--group-add", str(grp.getgrnam("render").gr_gid),
    ],
    "nvidia": [
        "--gpus", "all",
    ],
}

DOCKER_IMAGE_BY_GPU = {
    "amd": "ghcr.io/ggml-org/llama.cpp:full-vulkan",
    "intel": "ghcr.io/ggml-org/llama.cpp:full-intel",
    "nvidia": "ghcr.io/ggml-org/llama.cpp:full-cuda",
}

MODELS = [
    {
        "paths": ["qwen3.5/Qwen3.5-2B-UD-Q4_K_XL.gguf"],
        "full_attn_layers": 6, "kv_heads": 2, "head_dim": 256,
        "total_layers": 28,
        "moe": False,
        "urls": [
            "https://huggingface.co/unsloth"
            "/Qwen3.5-2B-GGUF/resolve/main/Qwen3.5-2B-UD-Q4_K_XL.gguf"
            "?download=true",
        ],
    },
    {
        "paths": ["qwen3.5/Qwen3.5-4B-UD-Q4_K_XL.gguf"],
        "full_attn_layers": 8, "kv_heads": 4, "head_dim": 256,
        "total_layers": 32,
        "moe": False,
        "urls": [
            "https://huggingface.co/unsloth"
            "/Qwen3.5-4B-GGUF/resolve/main/Qwen3.5-4B-UD-Q4_K_XL.gguf"
            "?download=true",
        ],
    },
    {
        "paths": ["qwen3.5/Qwen3.5-9B-Q4_K_M.gguf"],
        "full_attn_layers": 8, "kv_heads": 4, "head_dim": 256,
        "total_layers": 40,
        "moe": False,
        "urls": [
            "https://huggingface.co/unsloth"
            "/Qwen3.5-9B-GGUF/resolve/main/Qwen3.5-9B-Q4_K_M.gguf"
            "?download=true",
        ],
    },
    {
        "paths": ["qwen3.5/Qwen3.5-9B-UD-Q4_K_XL.gguf"],
        "full_attn_layers": 8, "kv_heads": 4, "head_dim": 256,
        "total_layers": 40,
        "moe": False,
        "urls": [
            "https://huggingface.co/unsloth"
            "/Qwen3.5-9B-GGUF/resolve/main/Qwen3.5-9B-UD-Q4_K_XL.gguf"
            "?download=true",
        ],
    },
    {
        "paths": ["qwen3.5/Qwen3.5-27B-Q4_0.gguf"],
        "full_attn_layers": 16, "kv_heads": 4, "head_dim": 256,
        "total_layers": 64,
        "moe": False,
        "urls": [
            "https://huggingface.co/unsloth"
            "/Qwen3.5-27B-GGUF/resolve/main/Qwen3.5-27B-Q4_0.gguf"
            "?download=true",
        ],
    },
    {
        "paths": ["qwen3.5/Qwen3.5-27B-UD-Q4_K_XL.gguf"],
        "full_attn_layers": 16, "kv_heads": 4, "head_dim": 256,
        "total_layers": 64,
        "moe": False,
        "urls": [
            "https://huggingface.co/unsloth"
            "/Qwen3.5-27B-GGUF/resolve/main/Qwen3.5-27B-UD-Q4_K_XL.gguf"
            "?download=true",
        ],
    },
    {
        "paths": ["qwen3.5/Qwen3.5-35B-A3B-Q4_K_M.gguf"],
        "full_attn_layers": 10, "kv_heads": 2, "head_dim": 256,
        "total_layers": 64,
        "moe": True,
        "urls": [
            "https://huggingface.co/unsloth"
            "/Qwen3.5-35B-A3B-GGUF/resolve/main/Qwen3.5-35B-A3B-Q4_K_M.gguf"
            "?download=true",
        ],
    },
    {
        "paths": ["qwen3.5/Qwen3.5-35B-A3B-UD-Q4_K_XL.gguf"],
        "full_attn_layers": 10, "kv_heads": 2, "head_dim": 256,
        "total_layers": 64,
        "moe": True,
        "urls": [
            "https://huggingface.co/unsloth"
            "/Qwen3.5-35B-A3B-GGUF/resolve/main"
            "/Qwen3.5-35B-A3B-UD-Q4_K_XL.gguf"
            "?download=true",
        ],
    },
    {
        "paths": ["qwen3.5/Qwen3.5-122B-A10B-UD-Q2_K_XL.gguf"],
        "full_attn_layers": 12, "kv_heads": 2, "head_dim": 256,
        "total_layers": 80,
        "moe": True,
        "urls": [
            "https://huggingface.co/unsloth"
            "/Qwen3.5-122B-A10B-GGUF/resolve/main"
            "/Qwen3.5-122B-A10B-UD-Q2_K_XL.gguf"
            "?download=true",
        ],
    },
    {
        "paths": [
            "qwen3.5/UD-Q4_K_XL"
            "/Qwen3.5-122B-A10B-UD-Q4_K_XL-00001-of-00003.gguf",
            "qwen3.5/UD-Q4_K_XL"
            "/Qwen3.5-122B-A10B-UD-Q4_K_XL-00002-of-00003.gguf",
            "qwen3.5/UD-Q4_K_XL"
            "/Qwen3.5-122B-A10B-UD-Q4_K_XL-00003-of-00003.gguf",
        ],
        "full_attn_layers": 12, "kv_heads": 2, "head_dim": 256,
        "total_layers": 80,
        "moe": True,
        "urls": [
            "https://huggingface.co/unsloth"
            "/Qwen3.5-122B-A10B-GGUF/resolve/main/UD-Q4_K_XL"
            "/Qwen3.5-122B-A10B-UD-Q4_K_XL-00001-of-00003.gguf"
            "?download=true",
            "https://huggingface.co/unsloth"
            "/Qwen3.5-122B-A10B-GGUF/resolve/main/UD-Q4_K_XL"
            "/Qwen3.5-122B-A10B-UD-Q4_K_XL-00002-of-00003.gguf"
            "?download=true",
            "https://huggingface.co/unsloth"
            "/Qwen3.5-122B-A10B-GGUF/resolve/main/UD-Q4_K_XL"
            "/Qwen3.5-122B-A10B-UD-Q4_K_XL-00003-of-00003.gguf"
            "?download=true",
        ],
    },
]

CONTEXT_SIZES = [
    1024,
    2048,
    4096,
    8192,
    16384,
    32768,
    65536,
    131072,
    196608,
    262144
]

DOCKER_ARGS_BEFORE = [
    "docker", "run", "--rm", "--entrypoint", "/app/llama-bench",
]
DOCKER_ARGS_AFTER = [
    "-t", str(psutil.cpu_count(logical=False) or 1),
    "--n-gen", "512",
    "--cache-type-k", "q8_0", "--cache-type-v", "q8_0",
    "--flash-attn", "on",
    "-r", "1",
]


def get_model_size_est(model_dir, model):
    total = 0
    for path in model["paths"]:
        total += (pathlib.Path(model_dir) / path).stat().st_size
    return total


def estimate_kv_cache_size(model, ctx_size, bytes_per_element=1):
    """
    q8 KV cache estimate for Qwen3.5. Only the full-attention ("Gated
    Attention") layers scale with ctx_size; the 3-in-4 Gated DeltaNet
    (linear attention) layers hold a fixed-size state, not included here.
    Values from each checkpoint's config.json text_config.
    """
    per_token_bytes = (
        2
        * model["full_attn_layers"]
        * model["kv_heads"]
        * model["head_dim"]
        * bytes_per_element
    )
    return per_token_bytes * ctx_size


def estimate_model_mem_usage(model_dir, model, ctx_size):
    return (
        get_model_size_est(model_dir, model)
        + estimate_kv_cache_size(model, ctx_size)
    )


def output(fd, out):
    print(out)
    fd.write(out + "\n")
    fd.flush()


def _compute_partial_offload(log_fd, model_dir, model, ctx, system):
    """Callback for scenarios that need dynamic -ngl calculation.
    """
    single_vram = system.gpu["single_vram"]
    model_size = get_model_size_est(model_dir, model)
    kv_cache = estimate_kv_cache_size(model, ctx)
    total_size = 2 * (model_size + kv_cache)

    # Check if full offload fits
    if total_size <= single_vram:
        output(log_fd, "full offload")
        return ["-ngl", "999"]

    # Calculate partial offload
    total_layers = model.get("total_layers", 0)
    if total_layers == 0:
        output(log_fd, "no total_layers defined, cannot calculate partial offload")
        return None

    per_layer_mem = total_size / total_layers
    available_for_layers = single_vram

    if available_for_layers <= 0:
        output(log_fd, "not enough VRAM")
        return None

    ngl = int(available_for_layers // per_layer_mem)

    if ngl < 1:
        output(
            log_fd,
            "model too large for partial offload"
            f" (need ~{per_layer_mem * total_layers // 1024 // 1024}MB,"
            f" have ~{single_vram // 1024 // 1024}MB)"
        )
        return None

    output(log_fd, f"partial offload: {ngl}/{total_layers} layers")
    return ["-ngl", str(ngl)]


SCENARIOS = [
    {
        "name": "GPU only",
        "args": {
            "docker": [],
            "llama-bench": ["-ngl", "999"],
        },
        "max_context": 9999999999999999999,
        "max_size": "get_vram",
        "only_moe": False,
    },
    {
        "name": "1 GPU + --cpu-moe",
        "args": {
            "docker": ["--env", "CUDA_VISIBLE_DEVICES=0"],
            "llama-bench": ["--n-cpu-moe", "999"],
        },
        "max_context": 9999999999999999999,
        "max_size": "get_ram",
        "only_moe": True,
    },
    {
        "name": "1 GPU + CPU offload",
        "args": {
            "docker": ["--env", "CUDA_VISIBLE_DEVICES=0"],
            "llama-bench": [],
        },
        "max_context": 9999999999999999999,
        "max_size": "get_ram",
        "only_moe": False,
        "extra_args": _compute_partial_offload,
    },
    {
        "name": "full CPU",
        "args": {
            "docker": ["--env", "CUDA_VISIBLE_DEVICES=-1"],
            "llama-bench": ["-ngl", "0"],
        },
        "max_context": 32768,
        "max_size": "get_ram",
        "only_moe": False,
        "skip_gpu_devices": True,
    },
]


def _run(cmd):
    try:
        return (
            subprocess.check_output(cmd, stderr=subprocess.DEVNULL)
            .decode().strip()
        )
    except (subprocess.CalledProcessError, FileNotFoundError):
        return ""


def get_gpus():
    result = {
        "amd": {"count": 0, "single_vram": 0, "total_vram": 0},
        "intel": {"count": 0, "single_vram": 0, "total_vram": 0},
        "nvidia": {"count": 0, "single_vram": 0, "total_vram": 0},
    }

    # NVIDIA via nvidia-smi
    nvidia_output = _run([
        "nvidia-smi",
        "--query-gpu=memory.total",
        "--format=csv,noheader,nounits"
    ])
    if nvidia_output:
        for line in nvidia_output.splitlines():
            match = re.search(r"(\d+)", line)
            if match is None:
                continue
            mb = int(match.group(1))
            result["nvidia"]["count"] += 1
            if result["nvidia"]["single_vram"] == 0:
                result["nvidia"]["single_vram"] = (mb * 1024 * 1024)
            result["nvidia"]["total_vram"] += (mb * 1024 * 1024)

    # AMD via rocm-smi
    amd_output = _run(["rocm-smi", "--showmeminfo", "vram"])
    if amd_output:
        for line in amd_output.splitlines():
            if "Total Memory" not in line:
                continue
            m = re.search(r": (\d+)", line)
            if m:
                # rocm-smi reports in bytes
                result["amd"]["count"] += 1
                if result["amd"]["single_vram"] == 0:
                    result["amd"]["single_vram"] = int(m.group(1))
                result["amd"]["total_vram"] += int(m.group(1))

    # Intel via sysfs
    # Deduplicate Intel (card0, card1 can point to same device)
    devices = set()
    for entry in os.listdir("/sys/class/drm"):
        if entry.startswith("card") and os.path.islink(f"/sys/class/drm/{entry}/device"):
            dev = os.path.realpath(f"/sys/class/drm/{entry}/device")
            devices.add(dev)
    for device in devices:
        vendor_file = os.path.join(device, "device", "vendor")
        if os.path.exists(vendor_file):
            with open(vendor_file) as vendor_fd:
                vendor_id = vendor_fd.read().strip()
            if vendor_id == "0x8086":  # Intel
                # TODO(Jflesch): Assuming B60 here.
                total_bytes = 24 * (1024 ** 3)
                result["intel"]["count"] += 1
                if result["intel"]["single_vram"] == 0:
                    result["intel"]["single_vram"] = total_bytes
                result["intel"]["total_vram"] += total_bytes

    return result


class System:
    def __init__(self):
        self.hostname = socket.gethostname()
        self.ram = psutil.virtual_memory().total

        self.gpus = get_gpus()

        # pick the GPU type with the most total VRAM
        best_gpu_type = max(
            (gpu[1]["total_vram"], gpu[0])
            for gpu in self.gpus.items()
        )[1]

        self.gpu_type = best_gpu_type
        self.gpu = self.gpus[best_gpu_type]

    def get_ram(self):
        return self.ram

    def get_vram(self):
        return self.gpu["total_vram"]


def download(url, file_path):
    print(f"{url} -> {file_path} …")

    file_dir = os.path.dirname(file_path)
    os.makedirs(file_dir, exist_ok=True)
    last_display = 0
    downloaded = 0
    with requests.get(url, stream=True) as response:
        response.raise_for_status()
        total = int(response.headers.get("content-length", 1))
        with open(file_path + ".part", "wb") as out:
            for chunk in response.iter_content(chunk_size=8192):
                out.write(chunk)
                downloaded += len(chunk)
                if last_display == 0 or last_display + (10 * (1024 ** 2)) < downloaded:
                    last_display = downloaded
                    sys.stdout.write(
                        "\33[2K\r"
                        f"{downloaded // (1024 ** 2)} MB"
                        f" / {total // (1024 ** 2)} MB"
                    )
                    sys.stdout.flush()
            print(
                "\33[2K\r"
                f"{downloaded // (1024 ** 2)} MB"
                f" / {total // (1024 ** 2)} MB"
            )

    os.rename(file_path + ".part", file_path)


def main():
    if len(sys.argv) <= 2:
        print("Usage:")
        print(f"  {sys.argv[0]} <model storage path> <output file>")
        print("Example:")
        print("   ./llm_benchmark.py /data/llm benchmark.txt")
        sys.exit(1)

    model_dir = sys.argv[1]

    with open(sys.argv[2], "w") as fd:
        system = System()
        output(fd, f"System: {system.hostname}")
        output(fd, f"RAM: {system.ram // 1024 // 1024} MB")
        output(fd, f"GPU type: {system.gpu_type}")
        output(fd, f"Number of GPUs: {system.gpu['count']}")
        output(
            fd,
            f"Single VRAM: {system.gpu['single_vram'] // 1024 // 1024} MB"
        )
        output(
            fd,
            f"Total VRAM: {system.gpu['total_vram'] // 1024 // 1024} MB"
        )
        output(fd, "")

        print(f"Checking all models are available in {model_dir} ...")
        for model in MODELS:
            model_path = os.path.join(model_dir, model["paths"][-1])
            sys.stdout.write(f"{model_path} … ")
            if os.path.exists(model_path):
                print("OK")
                continue
            print("KO. Downloading …")

            assert len(model["paths"]) == len(model["urls"])
            for (path, url) in zip(model["paths"], model["urls"]):
                gguf_path = os.path.join(model_dir, path)
                download(url, gguf_path)
            print("Done")

        for scenario in SCENARIOS:
            output(fd, "")
            output(fd, f"- Scenario: {scenario['name']}")

            max_size = getattr(system, scenario["max_size"])()
            max_size -= (1024 ** 3)
            output(fd, f"Max size = {max_size // 1024 // 1024} MB")

            for model in MODELS:
                if scenario["only_moe"] and not model["moe"]:
                    output(
                        fd,
                        f"scenario '{scenario['name']}' is for MoE only,"
                        f" but {model['paths']} is not a MoE. Skipped"
                    )
                    continue

                for ctx in CONTEXT_SIZES:
                    output(fd, "")
                    output(fd, f"--- Benchmark Started {time.ctime()} ---")
                    output(fd, f"scenario: {scenario['name']}")
                    output(fd, f"model: {model['paths']}")
                    output(fd, f"context size: {ctx}")
                    if ctx > scenario["max_context"]:
                        output(fd, "TOO SLOW. Skipped")
                        continue
                    mem_required = estimate_model_mem_usage(
                        model_dir, model, ctx
                    )
                    if mem_required > max_size:
                        output(fd, "TOO BIG. Skipped")
                        continue

                    # Run scenario-specific pre-benchmark callback if defined

                    llama_bench_args = list(scenario["args"]["llama-bench"])
                    extra_args = []

                    extra_args_func = scenario.get("extra_args")
                    if extra_args_func is not None:
                        extra_args = extra_args_func(
                            fd, model_dir, model, ctx, system
                        )
                        if extra_args is None:
                            output(fd, "SKIPPED")
                            continue

                    cmd = (
                        DOCKER_ARGS_BEFORE
                        + ["-v", f"{model_dir}:/models:ro"]
                        + scenario["args"]["docker"]
                        + ([] if scenario.get("skip_gpu_devices")
                           else DOCKER_DEVICE_ARGS_BY_GPU[system.gpu_type])
                        + [DOCKER_IMAGE_BY_GPU[system.gpu_type]]
                        + DOCKER_ARGS_AFTER
                        + extra_args
                        + llama_bench_args
                        + ["--n-prompt", str(ctx)]
                        + ["-m", f"/models/{model['paths'][0]}"]
                    )
                    output(fd, "Running:")
                    output(fd, " ".join(cmd))

                    try:
                        result = subprocess.check_output(
                            " ".join(cmd),
                            shell=True,
                            stderr=subprocess.STDOUT,
                            timeout=7200
                        )
                        output(fd, result.decode('utf-8'))
                    except subprocess.TimeoutExpired:
                        output(fd, "ERROR: Timeout (2h)")
                    except subprocess.CalledProcessError as e:
                        output(fd, f"ERROR: {e.output.decode('utf-8')}")

        output(fd, f"--- END: {time.ctime()} ---")


if __name__ == "__main__":
    main()