#!/usr/bin/env python3
"""SEE model autotest — проверяет все загруженные модели после обновления.

Универсальный: не содержит хардкода имён моделей, подходит для любой SEE-инсталляции.
Текст запроса подбирается динамически по имени модели.
Для точных проверок можно подключить asserts-файл.

Использование:
    python3 test_see_models.py
    python3 test_see_models.py --host 10.101.60.26
    python3 test_see_models.py --host 10.101.60.25 --port 6184
    python3 test_see_models.py --asserts see_asserts.json

Файл ассертов (JSON):
    {
      "asserts": [
        {
          "model": "fio",              // подстрока для поиска модели
          "text": "Иванов Иван Иванович", // текст запроса (опционально)
          "checks": {
            "min_found": 1,            // мин. количество сущностей/классов
            "text_contains": "иван",    // подстрока в text или class
            "calculated_contains": "Ф:", // подстрока в calculated
            "confidence_min": 0.5       // мин. confidence
          }
        }
      ]
    }

Возвращает exit code 0 если все модели прошли, иначе 1.
"""

import argparse
import json
import sys
import os
import unittest
from urllib.parse import quote

try:
    import urllib.request as urlreq
    import urllib.error as urlerr
except ImportError:
    urlreq = None

GREEN = '\033[92m'
RED = '\033[91m'
YELLOW = '\033[93m'
BOLD = '\033[1m'
RESET = '\033[0m'

def _disable_colors():
    global GREEN, RED, YELLOW, BOLD, RESET
    GREEN = RED = YELLOW = BOLD = RESET = ''

if not os.isatty(sys.stdout.fileno()) or os.environ.get('NO_COLOR'):
    _disable_colors()

DEFAULT_HOST = "10.101.60.25"
DEFAULT_PORT = 6184
BASE_PATH = "/see"

_DEFAULT_ARGS = argparse.Namespace(
    host=DEFAULT_HOST,
    port=DEFAULT_PORT,
    asserts=None,
    verbose=False,
    no_color=False
)

TEXT_RULES = [
    # (keyword, text) — первое совпадение побеждает
    ("fio", "Иванов Иван Иванович"),
    ("phone", "8 913 456 78 90"),
    ("number", "пятьсот двадцать три"),
    ("date", "на завтра"),
    ("record_date", "на завтра"),
    ("time", "в пять часов вечера"),
    ("record_time", "в пять часов вечера"),
    ("uuid", "550e8400-e29b-41d4-a716-446655440000"),
    ("zags", "загс центрального района"),
    ("nsk_fillials", "родники"),
    ("filial", "в филиал на родниках"),
    ("offices_mfc", "родники"),
    ("без уточнения возраста", "вписать ребенка в загранпаспорт"),
    ("тип загранпаспорта", "новый"),
    ("возраст", "мне 25 лет"),
    ("загран", "вписать ребенка в загранпаспорт"),
    ("yesno", "да"),
    ("weekday", "в пятницу"),
    ("дни недели", "в пятницу"),
    ("operator", "соедините с оператором"),
    ("end_repeat", "все, до свидания"),
    ("повтор", "повторите ещё раз"),
    ("other_payment", "мобилизационные выплаты"),
    ("payment", "мобилизационные выплаты"),
    ("mobilizac", "оператор"),
    ("trud", "зарегистрироваться в цзн"),
    ("subsidii", "субсидия на жкх"),
    ("minsoc", "субсидия на жкх"),
    ("matcap", "на улучшение жилищных условий"),
    ("маткап", "на улучшение жилищных условий"),
    ("паспорт", "загран"),
    ("пособие", "пособие по рождению и воспитанию детей"),
    ("начало диалог", "старт"),
    ("граждан", "я гражданин россии"),
    ("выплат", "пособие по рождению и воспитанию детей"),
    ("статус", "заявление зарегистрировано"),
    ("суперсервис", "рождение ребенка"),
    ("рождение", "рождение ребенка"),
    ("лаврентьев", "лаврентьевская карта"),
    ("несогласи", "это не то время"),
    ("отказ", "отказываюсь"),
    ("оценк", "пять"),
    ("grade", "пять"),
    ("rating", "пять"),
    ("исключение", "вопрос по жкх"),
    ("режим", "режим работы мфц"),
    ("жкх", "вопрос по жкх"),
    ("вернуться", "вернуться в главное меню"),
    ("не знаю", "не знаю код подразделения"),
    ("service_mfc", "хочу получить услугу"),
    ("service mfc", "хочу получить услугу"),
    ("мои", "мои талоны"),
    ("количество", "2"),
    ("mfc", "записаться в мфц"),
    ("мфц", "записаться в мфц"),
    ("тест", "тестовое сообщение"),
    ("test", "test message"),
    ("заглушк", "тестовый запрос"),
    ("stub", "test request"),
    ("конкретик", "конкретный вопрос"),
    ("да-нет", "да"),
]


def pick_text(model_id):
    name_lower = model_id.lower()
    for keyword, text in TEXT_RULES:
        if keyword in name_lower:
            return text
    return "Здравствуйте, хочу записаться на приём в МФЦ"


def load_asserts(path):
    if not path:
        return []
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    return data.get("asserts", [])


class SeeAPI:
    def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT):
        self.base = f"http://{host}:{port}{BASE_PATH}"

    def get_models(self):
        return json.loads(self._request(f"{self.base}/models"))

    def test_entity(self, model_id, text, similarity=70):
        encoded = quote(model_id, safe="")
        url = f"{self.base}/entities/{encoded}?text={quote(text)}&similarity={similarity}"
        return json.loads(self._request(url))

    def _request(self, url, timeout=30):
        try:
            with urlreq.urlopen(url, timeout=timeout) as r:
                return r.read().decode("utf-8")
        except urlerr.HTTPError as e:
            raise RuntimeError(f"HTTP {e.code}: {e.read().decode('utf-8')[:200]}")
        except urlerr.URLError as e:
            raise RuntimeError(f"Connection failed: {e.reason}")


class TestSeeModels(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        if not hasattr(cls, 'args'):
            cls.args = _DEFAULT_ARGS
        cls.api = SeeAPI(cls.args.host, cls.args.port)
        cls.asserts = load_asserts(cls.args.asserts)
        data = cls.api.get_models()
        TestSeeModels.models = data
        total = len(data.get("models", [])) + len(data.get("handlers", [])) + len(data.get("pretrained", []))
        print(f"\n  models: {len(data.get('models', []))}, "
              f"handlers: {len(data.get('handlers', []))}, "
              f"pretrained: {len(data.get('pretrained', []))}")
        if cls.asserts:
            print(f"  asserts loaded: {len(cls.asserts)}")
        if total == 0:
            raise unittest.SkipTest("Нет ни одной модели")

    def _find_assert(self, model_id):
        for a in self.asserts:
            if a["model"] == model_id:
                return a
        for a in self.asserts:
            if a["model"] in model_id:
                return a
        return None

    def _check_assert(self, model_id, result):
        a = self._find_assert(model_id)
        if not a:
            return
        checks = a.get("checks", {})
        has_classes = isinstance(result, dict) and "classes" in result
        if has_classes:
            classes = result.get("classes", [])
            if "min_found" in checks:
                self.assertGreaterEqual(len(classes), checks["min_found"],
                    f"{model_id}: expected >= {checks['min_found']} items, got {len(classes)}")
            if classes and "text_contains" in checks:
                self.assertIn(checks["text_contains"], classes[0].get("class", ""),
                    f"{model_id}: missing '{checks['text_contains']}'")
            if classes and "confidence_min" in checks:
                self.assertGreaterEqual(classes[0].get("confidence", 0), checks["confidence_min"],
                    f"{model_id}: confidence < {checks['confidence_min']}")
        else:
            entities = []
            for v in result.values():
                if isinstance(v, list):
                    entities.extend(v)
            count = len(entities)
            if "min_found" in checks:
                self.assertGreaterEqual(count, checks["min_found"],
                    f"{model_id}: expected >= {checks['min_found']} entities, got {count}")
            if count > 0:
                first = entities[0]
                if "text_contains" in checks:
                    self.assertIn(checks["text_contains"], first.get("text", ""),
                        f"{model_id}: text missing '{checks['text_contains']}'")
                if "calculated_contains" in checks:
                    self.assertIn(checks["calculated_contains"], first.get("calculated", ""),
                        f"{model_id}: calculated missing '{checks['calculated_contains']}'")
                if "confidence_min" in checks:
                    self.assertGreaterEqual(first.get("confidence", 0), checks["confidence_min"],
                        f"{model_id}: confidence < {checks['confidence_min']}")

    def _is_empty(self, result):
        if not result or result == {}:
            return True
        return all(
            isinstance(v, list) and len(v) == 0 for v in result.values()
        )

    def _fmt_result(self, result):
        if isinstance(result, dict) and "classes" in result:
            cs = result.get("classes", [])
            if cs:
                top = cs[0]
                return f"{top.get('class','?')} ({top.get('confidence',0):.3f})"
            return "∅"
        for entities in result.values():
            if isinstance(entities, list) and entities:
                e = entities[0]
                calc = e.get("calculated", "")
                if calc:
                    return calc[:60]
                return e.get("text", "?")[:40]
        return "∅"

    def _run_section(self, section_name, model_ids):
        """Тестирует список моделей, выводит компактно."""
        results = []
        for model_id in model_ids:
            if not model_id:
                continue
            text = pick_text(model_id)
            a = self._find_assert(model_id)
            if a and "text" in a:
                text = a["text"]
            try:
                result = self.api.test_entity(model_id, text)
                if self._is_empty(result):
                    results.append(("empty", model_id, None, None))
                else:
                    self._check_assert(model_id, result)
                    results.append(("ok", model_id, result, text))
            except AssertionError:
                results.append(("assert_fail", model_id, None, text))
                raise
            except Exception as e:
                results.append(("error", model_id, None, f"{e}"))

        passed = sum(1 for r in results if r[0] == "ok")
        empty = sum(1 for r in results if r[0] == "empty")
        errors = [r for r in results if r[0] == "error"]
        assert_fails = [r for r in results if r[0] == "assert_fail"]

        print(f"\n{BOLD}── {section_name} ──────────────────────────────────────────────{RESET}")
        if results:
            for status, mid, result, extra in results:
                if status == "ok" and self.args.verbose:
                    print(f"  {GREEN}✓ {mid}{RESET}  → {self._fmt_result(result)}")
                elif status == "ok":
                    print(f"  {GREEN}✓ {mid}{RESET}")
                elif status == "empty":
                    print(f"  {RED}✗ {mid}{RESET}  → empty")
                elif status == "assert_fail":
                    print(f"  {RED}✗ {mid}{RESET}  → assert fail")
                elif status == "error":
                    print(f"  {YELLOW}⚠ {mid}{RESET}  → {extra}")
        color = GREEN if empty == 0 and not errors and not assert_fails else RED
        print(f"  {BOLD}{color}─── {passed} passed, {empty} empty{RESET}", end="")
        if errors:
            print(f", {len(errors)} errors", end="")
        if assert_fails:
            print(f", {len(assert_fails)} assert fails", end="")
        print()

        error_msgs = []
        for r in errors:
            error_msgs.append(f"  {r[1]}: {r[3]}")
        if error_msgs:
            self.fail(f"Ошибки {section_name}:\n" + "\n".join(error_msgs))

    def test_01_list_models(self):
        """Получение списка моделей"""
        data = self.models
        self.assertIn("models", data)
        self.assertIn("handlers", data)
        self.assertIn("pretrained", data)

    def test_02_pretrained_models(self):
        """Проверка pretrained-моделей"""
        pretrained = self.models.get("pretrained", [])
        if not pretrained:
            self.skipTest("Нет pretrained моделей")
        self._run_section("PRETRAINED", pretrained)

    def test_03_handler_models(self):
        """Проверка handler-моделей"""
        handlers = self.models.get("handlers", [])
        if not handlers:
            self.skipTest("Нет handler моделей")
        self._run_section("HANDLERS", handlers)

    def test_04_custom_models(self):
        """Проверка кастомных моделей"""
        models = self.models.get("models", [])
        if not models:
            self.skipTest("Нет кастомных моделей")
        self._run_section("CUSTOM", models)


def main():
    parser = argparse.ArgumentParser(description="SEE model autotest")
    parser.add_argument("--host", default=DEFAULT_HOST, help=f"Хост (default: {DEFAULT_HOST})")
    parser.add_argument("--port", type=int, default=DEFAULT_PORT, help=f"Порт (default: {DEFAULT_PORT})")
    parser.add_argument("--asserts", default=None, help="JSON файл с ассертами")
    parser.add_argument("--verbose", "-v", action="store_true", help="Показать результат каждой модели")
    parser.add_argument("--no-color", action="store_true", help="Отключить цвета в выводе")
    args = parser.parse_args()
    if args.no_color:
        _disable_colors()
    TestSeeModels.args = args

    suite = unittest.TestLoader().loadTestsFromTestCase(TestSeeModels)
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(suite)
    sys.exit(0 if result.wasSuccessful() else 1)


if __name__ == "__main__":
    main()