"""
zmierz-model.py - szybkość, pamięć karty i koszt prądu lokalnego modelu (promptowy.com/lokalny-model-czy-api/)
Autor: Piotr Olszewski, promptowy.com. Licencja skryptu: MIT. Sprawdzone 26.09.2026: llama.cpp b11200, RTX 4090.

Działa z każdym serwerem zgodnym z API OpenAI: llama-server (port 8080), LM Studio (1234), Ollama (11434).
Wysyła 5 polskich poleceń, mierzy tokeny na sekundę, szczyt pamięci karty i średni pobór mocy (nvidia-smi),
a na końcu liczy koszt prądu za milion tokenów. Tylko biblioteka standardowa Pythona 3.9+.

Przykłady:
  python zmierz-model.py                                   (llama-server na 127.0.0.1:8080)
  python zmierz-model.py --url http://127.0.0.1:1234 --model gemma-4-12b      (LM Studio)
  python zmierz-model.py --url http://127.0.0.1:11434 --model gemma4:12b      (Ollama)
  python zmierz-model.py --cena 1.05 --reszta-w 100        (cena prądu zł/kWh, dopłata za resztę komputera)
"""
import argparse, json, statistics, subprocess, sys, threading, time, urllib.error, urllib.request

POLECENIA = [
    "Napisz uprzejmą odpowiedź na reklamację klienta, którego paczka spóźnia się 5 dni. Do 120 słów.",
    "Wyjaśnij dwunastolatkowi, czym jest inflacja. 100-130 słów, jeden przykład.",
    "Streść w 5 punktach: zarząd postanowił przesunąć premierę aplikacji o miesiąc, zwiększyć budżet testów o 20% i zatrudnić dwóch testerów.",
    "Zaproponuj 10 tytułów artykułu o oszczędzaniu prądu w bloku, każdy do 60 znaków.",
    "Przetłumacz na polski, bez kalk: 'We're over the moon to announce our new release - it's been a long time coming.'",
]


class Czujnik(threading.Thread):
    def __init__(self):
        super().__init__(daemon=True); self.stop = False; self.vram = []; self.moc = []
    def run(self):
        while not self.stop:
            try:
                o = subprocess.check_output(["nvidia-smi", "--query-gpu=memory.used,power.draw", "--format=csv,noheader,nounits"], text=True, timeout=5)
                v, w = [float(x) for x in o.strip().splitlines()[0].split(",")]
                self.vram.append(v); self.moc.append(w)
            except Exception:
                return
            time.sleep(0.5)


def main():
    a = argparse.ArgumentParser(description="Pomiar lokalnego modelu: tok/s, VRAM, koszt prądu")
    a.add_argument("--url", default="http://127.0.0.1:8080", help="adres serwera zgodnego z API OpenAI")
    a.add_argument("--model", default="lokalny", help="nazwa modelu (LM Studio i Ollama jej wymagają)")
    a.add_argument("--cena", type=float, default=1.05, help="cena prądu w zł za kWh (domyślnie 1,05 - G11 z dystrybucją, 2026)")
    a.add_argument("--reszta-w", type=float, default=0, help="dolicz tyle watów za resztę komputera (np. 100)")
    a.add_argument("--max-tokens", type=int, default=800)
    x = a.parse_args()
    try:
        o = subprocess.check_output(["nvidia-smi", "--query-gpu=name,memory.used", "--format=csv,noheader,nounits"], text=True, timeout=5)
        karta, tlo = o.strip().splitlines()[0].rsplit(",", 1); tlo = float(tlo)
        print(f"Karta: {karta.strip()}, zajęte przed testem: {tlo:.0f} MB")
    except Exception:
        karta, tlo = None, 0
        print("Brak nvidia-smi - mierzę tylko szybkość, bez pamięci i prądu.")
    # llama-server odpowiada 503, dopóki wczytuje model - czekamy do 5 minut
    for proba in range(150):
        try:
            urllib.request.urlopen(x.url.rstrip("/") + "/v1/models", timeout=5); break
        except urllib.error.HTTPError as e:
            if e.code != 503: break
        except Exception:
            pass
        if proba == 0: print("Czekam, aż serwer wczyta model...")
        time.sleep(2)
    wyniki = []
    for i, p in enumerate(POLECENIA, 1):
        body = {"model": x.model, "messages": [{"role": "user", "content": p}], "max_tokens": x.max_tokens, "temperature": 0.7}
        cz = Czujnik(); cz.start(); t0 = time.time()
        try:
            req = urllib.request.Request(x.url.rstrip("/") + "/v1/chat/completions", data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
            j = json.load(urllib.request.urlopen(req, timeout=900))
        except Exception as e:
            sys.exit(f"Serwer {x.url} nie odpowiada ({e}). Czy model jest uruchomiony?")
        czas = time.time() - t0; cz.stop = True; cz.join()
        tok = j.get("usage", {}).get("completion_tokens") or 0
        tps = j.get("timings", {}).get("predicted_per_second") or (tok / czas if czas else 0)
        moc = statistics.mean(cz.moc) if cz.moc else 0
        wyniki.append({"tok": tok, "czas": czas, "tps": tps, "vram": max(cz.vram or [0]), "moc": moc})
        tekst = (j["choices"][0]["message"].get("content") or "").strip().replace("\n", " ")
        print(f"[{i}/5] {tok} tokenów w {czas:.1f} s, {tps:.1f} tok/s, {moc:.0f} W | {tekst[:70]}...")
    tok = sum(w["tok"] for w in wyniki); czas = sum(w["czas"] for w in wyniki)
    print("\n=== Podsumowanie ===")
    print(f"Szybkość generowania: {statistics.mean(w['tps'] for w in wyniki):.1f} tok/s")
    if karta:
        wh = sum(w["czas"] * (w["moc"] + x.reszta_w) for w in wyniki) / 3600
        print(f"Pamięć karty: szczyt {max(w['vram'] for w in wyniki):.0f} MB, z czego model ok. {max(w['vram'] for w in wyniki) - tlo:.0f} MB")
        print(f"Energia: {wh:.2f} Wh na {tok} tokenów")
        if tok:
            print(f"Koszt prądu: {wh / 1000 * x.cena * 100:.3f} gr za test, {wh / 1000 * x.cena / tok * 1e6:.2f} zł za milion tokenów "
                  f"(cena {x.cena} zł/kWh{', +%.0f W reszta komputera' % x.reszta_w if x.reszta_w else ', sama karta'})")


if __name__ == "__main__":
    main()