#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Тестирование LLM-моделей для AI-ассистента «Зам» (АС «Платформа ОПОРА РОССИИ»).

Запуск (Windows PowerShell):
    $env:OPENROUTER_API_KEY = "sk-or-..."
    py -3 test_models.py

Свой список моделей:
    py -3 test_models.py --models "openai/gpt-4o-mini,anthropic/claude-3.5-haiku"

Результат: таблица в консоли + отчёт report_<дата>.md
"""

import argparse
import json
import os
import sys
import time
from datetime import datetime

import requests

OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
MODELS_URL = "https://openrouter.ai/api/v1/models"

DEFAULT_MODELS = [
    "openai/gpt-4o-mini",
    "anthropic/claude-3.5-haiku",
    "google/gemini-2.0-flash-001",
    "qwen/qwen-2.5-72b-instruct",
    "deepseek/deepseek-chat",
    "mistralai/mistral-small-24b-instruct-2501",
]

SYSTEM = (
    "Ты — ассистент «Зам» платформы ОПОРА РОССИИ. "
    "Отвечай по-русски, кратко и по делу. Термины: «точка 0», «текущий срез», "
    "«прирост», «динамика», «красная зона», «вклад округа»."
)

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "create_task",
            "description": "Создать задачу в трекере",
            "parameters": {
                "type": "object",
                "properties": {
                    "title": {"type": "string", "description": "Название задачи"},
                    "due_date": {"type": "string", "description": "Срок в формате YYYY-MM-DD"},
                    "priority": {"type": "string", "enum": ["низкий", "средний", "высокий"]},
                    "assignee": {"type": "string", "description": "Ответственный (регион или роль)"},
                },
                "required": ["title", "due_date", "priority", "assignee"],
            },
        },
    }
]

TESTS = [
    {
        "name": "ru_generation",
        "desc": "Русский текст: еженедельный фокус целей",
        "messages": [
            {"role": "system", "content": SYSTEM},
            {
                "role": "user",
                "content": (
                    "Составь сообщение «Еженедельный фокус целей» для Донецкой Народной "
                    "Республики. Прирост от точки 0: +136. Округ ЮФО: +365, с динамикой "
                    "12/12 регионов. Добавь рекомендацию «что делать»."
                ),
            },
        ],
    },
    {
        "name": "function_calling",
        "desc": "Вызов инструмента create_task",
        "messages": [
            {"role": "system", "content": SYSTEM},
            {
                "role": "user",
                "content": (
                    "Создай задачу: «Подготовить отчёт по приросту», срок 2026-10-01, "
                    "приоритет высокий, ответственный — регион ДНР."
                ),
            },
        ],
        "tools": TOOLS,
    },
    {
        "name": "json_output",
        "desc": "Структурированный JSON",
        "messages": [
            {"role": "system", "content": SYSTEM + " Отвечай строго в формате JSON, без пояснений."},
            {
                "role": "user",
                "content": (
                    "Верни JSON с полями region, point0, current, growth для региона ДНР: "
                    "точка 0 = 1000, текущий срез = 1136."
                ),
            },
        ],
    },
]


def load_pricing():
    """Возвращает {model_id: (prompt_price, completion_price)} в $ за 1 токен."""
    try:
        r = requests.get(MODELS_URL, timeout=30)
        r.raise_for_status()
        data = r.json().get("data", [])
        pricing = {}
        for m in data:
            p = m.get("pricing", {})
            try:
                pricing[m["id"]] = (float(p.get("prompt", 0)), float(p.get("completion", 0)))
            except (TypeError, ValueError):
                pass
        return pricing
    except Exception as e:
        print(f"[warn] не удалось загрузить прайс: {e}")
        return {}


def call_model(api_key, model, test):
    payload = {
        "model": model,
        "messages": test["messages"],
        "temperature": 0.3,
        "max_tokens": 700,
    }
    if "tools" in test:
        payload["tools"] = test["tools"]
        payload["tool_choice"] = "auto"

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://opora.my-dpr.ru",
        "X-Title": "OPORA Zam LLM test",
    }

    start = time.perf_counter()
    try:
        r = requests.post(OPENROUTER_URL, headers=headers, json=payload, timeout=120)
        latency = time.perf_counter() - start
        if r.status_code != 200:
            return {"ok": False, "error": f"HTTP {r.status_code}: {r.text[:300]}", "latency": latency}
        data = r.json()
    except Exception as e:
        return {"ok": False, "error": str(e), "latency": time.perf_counter() - start}

    choice = (data.get("choices") or [{}])[0]
    msg = choice.get("message", {})
    usage = data.get("usage", {}) or {}
    return {
        "ok": True,
        "latency": latency,
        "content": msg.get("content") or "",
        "tool_calls": msg.get("tool_calls") or [],
        "prompt_tokens": usage.get("prompt_tokens", 0),
        "completion_tokens": usage.get("completion_tokens", 0),
        "model_returned": data.get("model", model),
    }


def check_function_call(result):
    """Проверяет корректность вызова create_task."""
    if not result.get("ok"):
        return "ошибка"
    calls = result.get("tool_calls") or []
    if not calls:
        return "нет вызова"
    try:
        fn = calls[0]["function"]
        if fn["name"] != "create_task":
            return f"не та функция: {fn['name']}"
        args = json.loads(fn["arguments"])
        need = {"title", "due_date", "priority", "assignee"}
        missing = need - set(args)
        if missing:
            return f"нет полей: {', '.join(sorted(missing))}"
        if args.get("due_date") != "2026-10-01":
            return f"срок: {args.get('due_date')}"
        return "OK"
    except Exception as e:
        return f"ошибка разбора: {e}"


def check_json(result):
    if not result.get("ok"):
        return "ошибка"
    txt = (result.get("content") or "").strip()
    txt = txt.replace("```json", "").replace("```", "").strip()
    try:
        obj = json.loads(txt)
        need = {"region", "point0", "current", "growth"}
        missing = need - set(obj)
        if missing:
            return f"нет полей: {', '.join(sorted(missing))}"
        if int(obj.get("growth", 0)) != 136:
            return f"growth={obj.get('growth')}"
        return "OK"
    except Exception as e:
        return f"не JSON: {e}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", default=",".join(DEFAULT_MODELS))
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    api_key = os.environ.get("OPENROUTER_API_KEY")
    if not api_key:
        print("Ошибка: задайте переменную окружения OPENROUTER_API_KEY")
        sys.exit(1)

    models = [m.strip() for m in args.models.split(",") if m.strip()]
    pricing = load_pricing()

    rows = []
    for model in models:
        print(f"\n=== {model} ===")
        row = {"model": model}
        for test in TESTS:
            res = call_model(api_key, model, test)
            if not res["ok"]:
                print(f"  [{test['name']}] ОШИБКА: {res['error']}")
                row[test["name"]] = {"status": "ошибка", "latency": res["latency"]}
                continue

            if test["name"] == "function_calling":
                status = check_function_call(res)
            elif test["name"] == "json_output":
                status = check_json(res)
            else:
                status = "OK" if len(res["content"]) > 40 else "короткий ответ"

            pt, ct = res["prompt_tokens"], res["completion_tokens"]
            price = pricing.get(model, (0, 0))
            cost = pt * price[0] + ct * price[1]

            row[test["name"]] = {
                "status": status,
                "latency": res["latency"],
                "prompt_tokens": pt,
                "completion_tokens": ct,
                "cost_usd": cost,
                "sample": res["content"][:200],
            }
            print(f"  [{test['name']}] {status} | {res['latency']:.2f}s | "
                  f"{pt}+{ct} ток | ${cost:.6f}")

        rows.append(row)

    # Отчёт
    ts = datetime.now().strftime("%Y-%m-%d_%H%M")
    out = args.out or f"report_{ts}.md"
    lines = [f"# Тест LLM-моделей — {datetime.now():%Y-%m-%d %H:%M}", ""]
    lines.append("| Модель | ru_generation | function_calling | json_output | Ср. латентность | Стоимость теста |")
    lines.append("|---|---|---|---|---|---|")
    for row in rows:
        lat = [row[t["name"]]["latency"] for t in TESTS if row.get(t["name"], {}).get("latency")]
        avg_lat = sum(lat) / len(lat) if lat else 0
        cost = sum(row[t["name"]].get("cost_usd", 0) for t in TESTS if row.get(t["name"]))
        cells = []
        for t in TESTS:
            d = row.get(t["name"], {})
            cells.append(d.get("status", "—"))
        lines.append(f"| {row['model']} | " + " | ".join(cells) +
                     f" | {avg_lat:.2f}s | ${cost:.6f} |")
    lines.append("")
    lines.append("## Примеры ответов (ru_generation)")
    for row in rows:
        d = row.get("ru_generation", {})
        if d.get("sample"):
            lines.append(f"\n### {row['model']}\n\n{d['sample']}")

    with open(out, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))
    print(f"\nОтчёт сохранён: {out}")


if __name__ == "__main__":
    main()