# Hermes Dispatch Tool — LLM VRAM Calculator
# Source: https://hermesdispatch.dev/tools/vram-calculator/
# Description: Estimate GPU memory required to run a large language model locally.
# License: MIT (generated by hermesdispatch.dev)
#
# INSTALL:
#   1. Save this file as ~/.hermes/hermes-agent/tools/llm_vram_calculator.py
#   2. Restart Hermes or run /reset in a session
#   3. The tool auto-registers if Hermes uses auto-discovery of tools/*.py
#
# MANUAL REGISTRY (if auto-discovery is off):
#   from tools.llm_vram_calculator import register
#   register()

import json

DATA = {}

def _ok(result):
    return json.dumps({"success": True, "data": result}, indent=2)

def _err(message):
    return json.dumps({"success": False, "error": message}, indent=2)


TOOL_NAME = "llm_vram_calculator"
TOOLSET = "hardware"

SCHEMA = {
  "type": "function",
  "function": {
    "name": "llm_vram_calculator",
    "description": "Estimate GPU VRAM for a local LLM given parameters, quantization bits, context length, and optional KV-cache overhead.",
    "parameters": {
      "type": "object",
      "properties": {
        "parameters_b": {
          "type": "number",
          "description": "Model parameter count in billions."
        },
        "quantization_bits": {
          "type": "number",
          "description": "Bits per parameter."
        },
        "context_length": {
          "type": "integer",
          "default": 4096,
          "description": "Typical context length in tokens."
        },
        "include_kv_cache": {
          "type": "boolean",
          "default": true,
          "description": "Include 20% KV-cache/overhead buffer."
        }
      },
      "required": [
        "parameters_b",
        "quantization_bits"
      ]
    }
  }
}

def _run(args):
    parameters_b = float(args.get("parameters_b", 8))
    quantization_bits = float(args.get("quantization_bits", 4))
    context_length = int(args.get("context_length", 4096))
    include_kv_cache = bool(args.get("include_kv_cache", True))
    overhead = 1.2 if include_kv_cache else 1.05
    base_gb = (parameters_b * quantization_bits / 8) * overhead
    context_cost = (context_length * parameters_b / 1_000_000) * 0.5 if include_kv_cache else 0
    total = base_gb + context_cost
    if total <= 12:
        recommendation = "RTX 4060 Ti 16GB, RX 7600 XT 16GB, or Apple M4 16GB."
    elif total <= 24:
        recommendation = "RTX 3090/4090 24GB, RX 7900 XT 20GB, or Apple M4 Max 36GB."
    elif total <= 48:
        recommendation = "RTX 3090 pair (48GB), RTX 4090 pair (48GB), or Mac Studio M2 Ultra 64GB."
    elif total <= 80:
        recommendation = "Mac Studio M2 Ultra 64GB, RTX 4090 x3 (72GB), or 48GB datacenter GPU pair."
    elif total <= 140:
        recommendation = "Mac Studio M3 Max 128GB, two 48GB datacenter GPUs, or RTX 4090 x6."
    else:
        recommendation = "Multi-GPU workstation or cloud inference."
    return _ok({
        "estimated_vram_gb": round(total, 2),
        "base_vram_gb": round(base_gb, 2),
        "context_cost_gb": round(context_cost, 2),
        "formula": f"{parameters_b}B x {quantization_bits}-bit / 8 x {overhead:.1f} + {context_cost:.1f} GB context",
        "recommended_hardware": recommendation,
    })

def HANDLER(args):
    try:
        return _run(args)
    except Exception as e:
        return _err(str(e))


def register():
    """Manual registry hook. Import and call this to register with Hermes."""
    try:
        from tools.registry import registry
        registry.register(
            name=TOOL_NAME,
            toolset=TOOLSET,
            schema=SCHEMA,
            handler=HANDLER,
        )
    except ImportError:
        print("Hermes registry not found; skipping manual registration.")

if __name__ == "__main__":
    # CLI smoke test
    print(HANDLER({}))