# -*- coding: utf-8 -*-
"""Честный бенчмарк ZuperAI: скорость (ток/сек) + сколько задач решено с самопроверкой.

Прогоняет набор задач через полный цикл (генерация кода + тесты судьи + прогон),
считает, сколько решено с первого раза. Плюс замер скорости генерации.

Запуск:
  venv/Scripts/python.exe bench.py [deepseek|lmstudio|local]

Результат можно вставить в README как честный бенчмарк (в отличие от маркетинговых цифр).
"""
import sys
import tempfile
import time

from webui import (deepseek_available, deepseek_generate,
                   lmstudio_available, lmstudio_generate, _strip_fences)
from verify_util import TEST_SYSTEM, run_tests, ensure_unittest_header, tail

TASKS = [
    "Функция, которая возвращает сумму двух чисел",
    "Функция, которая возвращает среднее арифметическое списка чисел",
    "Функция is_palindrome: возвращает True, если строка — палиндром",
    "Функция, которая возвращает количество слов в строке",
    "Функция, которая возвращает n-е число Фибоначчи",
    "Функция, которая разворачивает строку задом наперёд",
    "Функция, которая находит максимальный элемент списка",
]


def _gen(backend, msgs, max_tokens):
    if backend == "deepseek":
        return _strip_fences(deepseek_generate(msgs, max_tokens=max_tokens))
    if backend == "lmstudio":
        return _strip_fences(lmstudio_generate(msgs, max_tokens=max_tokens))
    from infer_smart import load, gen
    if not hasattr(_gen, "_m"):
        _gen._m = load()
    tok, model = _gen._m
    return _strip_fences(gen(tok, model, msgs, max_tokens=max_tokens))


def bench_speed(backend):
    msgs = [{"role": "user", "content": "Напиши функцию, которая сортирует список чисел"}]
    t0 = time.time()
    out = _gen(backend, msgs, 256)
    dt = time.time() - t0
    toks = max(1, len(out.split()))
    return toks, dt, toks / dt


def main():
    backend = sys.argv[1] if len(sys.argv) > 1 else ("deepseek" if deepseek_available() else "lmstudio")
    judge = "deepseek" if deepseek_available() else backend

    toks, dt, tps = bench_speed(backend)
    print(f"=== скорость ({backend}): {toks} токов за {dt:.1f}с ≈ {tps:.1f} ток/с ===\n")

    passed = 0
    total_t = time.time()
    for i, task in enumerate(TASKS, 1):
        sandbox = tempfile.mkdtemp(prefix="zuperai_bench_")
        t_msgs = [{"role": "system", "content": TEST_SYSTEM},
                  {"role": "user",
                   "content": f"Задача: {task}\n\nНапиши unittest-тесты, импортирующие "
                              f"решение из solution. Верни ТОЛЬКО код тестов."}]
        tests = ensure_unittest_header(_gen(judge, t_msgs, 1024))
        msgs = [{"role": "user",
                 "content": f"Задача: {task}\n\nТесты:\n{tests}\n\n"
                            f"Напиши solution.py, проходящий все тесты. Верни ТОЛЬКО код."}]
        code = _gen(backend, msgs, 1024)
        ok, out = run_tests(code, tests, sandbox, 30)
        print(f"  {i}/{len(TASKS)} {task[:45]:47s} -> {'OK' if ok else 'FAIL'}")
        passed += 1 if ok else 0

    total = time.time() - total_t
    print(f"\nИТОГ: {passed}/{len(TASKS)} задач решено с первого раза "
          f"({passed * 100 // len(TASKS)}%), время {total:.0f}с")


if __name__ == "__main__":
    main()