"""System profiler (PRIVATE to Hardware_Module): CPU, RAM (cgroup aware), GPUs, model size."""
from __future__ import annotations

import json
import os
import re
from pathlib import Path
from typing import Dict, Optional

GIB = 1024 ** 3


def _read_int(path: str) -> Optional[int]:
    try:
        raw = Path(path).read_text().strip()
        if raw == "max":
            return None
        value = int(raw)
        return value if value < (1 << 60) else None
    except Exception:
        return None


def cgroup_memory() -> Dict[str, Optional[int]]:
    """Container limits (GitHub Codespaces / Docker / Kaggle) - psutil alone over-reports."""
    limit = _read_int("/sys/fs/cgroup/memory.max") or _read_int("/sys/fs/cgroup/memory/memory.limit_in_bytes")
    used = _read_int("/sys/fs/cgroup/memory.current") or _read_int("/sys/fs/cgroup/memory/memory.usage_in_bytes")
    return {"limit": limit, "used": used}


def estimate_model(model_path: str, token: Optional[str] = None) -> Dict:
    """Parameter count (in billions) + layer geometry from config.json (local dir or hub repo)."""
    cfg: Dict = {}
    local = Path(str(model_path)) / "config.json"
    try:
        if local.exists():
            cfg = json.loads(local.read_text(encoding="utf-8"))
        else:
            from huggingface_hub import hf_hub_download
            cfg = json.loads(Path(hf_hub_download(str(model_path), "config.json", token=token or None)).read_text())
    except Exception:
        cfg = {}
    hidden = int(cfg.get("hidden_size", 0))
    layers = int(cfg.get("num_hidden_layers", 0))
    vocab = int(cfg.get("vocab_size", 0))
    inter = int(cfg.get("intermediate_size", 4 * hidden if hidden else 0))
    heads = int(cfg.get("num_attention_heads", 1) or 1)
    kv_heads = int(cfg.get("num_key_value_heads", heads) or heads)
    head_dim = int(cfg.get("head_dim", hidden // heads if hidden else 0) or 0)
    tied = bool(cfg.get("tie_word_embeddings", False))
    if hidden and layers and vocab:
        attn = hidden * heads * head_dim + 2 * hidden * kv_heads * head_dim + heads * head_dim * hidden
        mlp = 3 * hidden * inter
        per_layer = attn + mlp + 2 * hidden
        embed = vocab * hidden
        total = embed * (1 if tied else 2) + layers * per_layer + hidden
    else:  # fall back to the name (e.g. "Qwen2.5-0.5B-Instruct", "...-9B")
        match = re.search(r"(\d+(?:\.\d+)?)\s*[bB]", str(model_path))
        total = int(float(match.group(1)) * 1e9) if match else int(1e9)
        layers = layers or 24
        embed = int(total * 0.1)
        per_layer = (total - embed) // layers
    return {"params": int(total), "params_b": round(total / 1e9, 3), "num_layers": layers,
            "layer_params": int(per_layer), "embed_params": int(embed), "tied_embeddings": tied,
            "hidden_size": hidden, "vocab_size": vocab}


def profile(model_path: str, token: Optional[str] = None) -> Dict:
    import psutil
    vm = psutil.virtual_memory()
    cg = cgroup_memory()
    total = vm.total
    avail = vm.available
    if cg["limit"]:
        total = min(total, cg["limit"])
        if cg["used"] is not None:
            avail = min(avail, max(cg["limit"] - cg["used"], 0))
    gpus = []
    try:
        import torch
        if torch.cuda.is_available():
            for i in range(torch.cuda.device_count()):
                free, tot = torch.cuda.mem_get_info(i)
                major, _minor = torch.cuda.get_device_capability(i)
                gpus.append({"index": i, "name": torch.cuda.get_device_name(i), "free_bytes": int(free),
                             "total_bytes": int(tot), "bf16": major >= 8})
    except Exception:
        gpus = []
    model = estimate_model(model_path, token)
    return {
        "cpu_cores": os.cpu_count() or 1,
        "ram_total_bytes": int(total), "ram_avail_bytes": int(avail),
        "gpu_count": len(gpus), "gpus": gpus,
        "gpu_free_bytes": sum(g["free_bytes"] for g in gpus),
        "model": model,
    }


def describe(prof: Dict) -> str:
    m = prof["model"]
    lines = [
        f"   CPU cores      : {prof['cpu_cores']}",
        f"   RAM total/free : {prof['ram_total_bytes'] / GIB:.2f} GiB / {prof['ram_avail_bytes'] / GIB:.2f} GiB",
        f"   GPUs           : {prof['gpu_count']}"
        + "".join(f"\n      cuda:{g['index']} {g['name']} free {g['free_bytes'] / GIB:.1f}/{g['total_bytes'] / GIB:.1f} GiB"
                  for g in prof["gpus"]),
        f"   Model          : ~{m['params_b']}B params, {m['num_layers']} layers, hidden {m['hidden_size']}, vocab {m['vocab_size']}",
    ]
    return "\n".join(lines)