"""
zmierz-model.py - velocidade, memória da placa e custo de energia de um modelo local
Autor: Piotr Olszewski. Licença do script: MIT. Testado em 26/09/2026: llama.cpp b11200, RTX 4090.

Funciona com qualquer servidor compatível com a API da OpenAI: llama-server (porta 8080), LM Studio (1234), Ollama (11434).
Envia cinco instruções traduzidas do polonês, mede tokens por segundo, pico de memória da placa e consumo médio (nvidia-smi),
e calcula o custo de energia por milhão de tokens. Usa apenas a biblioteca padrão do Python 3.9 ou posterior.

Exemplos:
  python zmierz-model.py                                   (llama-server em 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 0.2722 --reszta-w 100      (custo de energia em US$/kWh, adicional para o restante do computador)
"""
import argparse, json, statistics, subprocess, sys, threading, time, urllib.error, urllib.request

POLECENIA = [
    "Tradução da instrução original em polonês: Escreva uma resposta educada à reclamação de um cliente cuja encomenda está cinco dias atrasada. Até 120 palavras.",
    "Tradução da instrução original em polonês: Explique a uma criança de 12 anos o que é inflação. De 100 a 130 palavras, com um exemplo.",
    "Tradução da instrução original em polonês: Resuma em cinco tópicos: a diretoria decidiu adiar o lançamento do aplicativo por um mês, aumentar o orçamento de testes em 20% e contratar dois testadores.",
    "Tradução da instrução original em polonês: Sugira 10 títulos para um artigo sobre economia de energia em um apartamento, cada um com até 60 caracteres.",
    "Tradução da instrução original em polonês: Traduza para o polonês, sem traduções literais: 'Estamos muito felizes em anunciar nosso novo lançamento - ele demorou bastante para chegar.'",
]


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="Medição de modelo local: tokens/s, VRAM, custo de energia")
    a.add_argument("--url", default="http://127.0.0.1:8080", help="endereço do servidor compatível com a API da OpenAI")
    a.add_argument("--model", default="local", help="nome do modelo (obrigatório no LM Studio e no Ollama)")
    a.add_argument("--cena", type=float, default=0.2722, help="custo de energia em US$ por kWh (padrão: 0,2722)")
    a.add_argument("--reszta-w", type=float, default=0, help="adicione esta quantidade de watts para o restante do computador (por exemplo, 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"Placa: {karta.strip()}, memória ocupada antes do teste: {tlo:.0f} MB")
    except Exception:
        karta, tlo = None, 0
        print("nvidia-smi não encontrado - medindo apenas a velocidade, sem memória e consumo de energia.")
    # llama-server retorna 503 enquanto carrega o modelo - aguarde até cinco minutos
    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("Aguardando o servidor carregar o modelo...")
        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"O servidor {x.url} não responde ({e}). O modelo está em execução?")
        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} tokens em {czas:.1f} s, {tps:.1f} tokens/s, {moc:.0f} W | {tekst[:70]}...")
    tok = sum(w["tok"] for w in wyniki); czas = sum(w["czas"] for w in wyniki)
    print("\n=== Resumo ===")
    print(f"Velocidade média de geração: {statistics.mean(w['tps'] for w in wyniki):.1f} tokens/s")
    if karta:
        wh = sum(w["czas"] * (w["moc"] + x.reszta_w) for w in wyniki) / 3600
        print(f"Memória da placa: pico de {max(w['vram'] for w in wyniki):.0f} MB, dos quais cerca de {max(w['vram'] for w in wyniki) - tlo:.0f} MB para o modelo")
        print(f"Energia: {wh:.2f} Wh para {tok} tokens")
        if tok:
            print(f"Custo de energia: US$ {wh / 1000 * x.cena * 100:.3f} centavos no teste, US$ {wh / 1000 * x.cena / tok * 1e6:.2f} por milhão de tokens "
                  f"(US$ {x.cena}/kWh{', mais %.0f W para o restante do computador' % x.reszta_w if x.reszta_w else ', apenas a placa'})")


if __name__ == "__main__":
    main()