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

Использование:
    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",               // подстрока для поиска модели
          "type": "entity|classifier",
          "text": "Иванов Иван Иванович", // текст для запроса
          "checks": {
            // для entity:
            "min_found": 1,              // мин. количество найденных сущностей
            "text_contains": "иванов",    // подстрока в text первой сущности
            "calculated_contains": "Ф:",  // подстрока в calculated
            "confidence_min": 0.5,        // мин. confidence

            // для classifier:
            "min_classes": 1,            // мин. количество классов
            "top_class_contains": "Yes",  // подстрока в top классе
            "confidence_min": 0.5         // мин. confidence топ-класса
          }
        }
      ]
    }

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

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

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

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

PRETRAINED_TESTS = {
    "fio": {"text": "Иванов Иван Иванович", "hint": "ФИО"},
    "phone": {"text": "8 913 456 78 90", "hint": "телефон"},
    "number": {"text": "пятьсот двадцать три", "hint": "число"},
    "date": {"text": "на завтра", "hint": "дата"},
    "uuid": {"text": "550e8400-e29b-41d4-a716-446655440000", "hint": "UUID"},
    "time": {"text": "в пять часов вечера", "hint": "время"},
}

HANDLER_TESTS = {
    "fio2": {"text": "Иванов Иван Иванович", "type": "entity"},
    "phone2": {"text": "8 913 456 78 90", "type": "entity"},
    "record_date": {"text": "на завтра", "type": "entity"},
    "record_time": {"text": "в пять часов вечера", "type": "entity"},
    "yesno": {"text": "да", "type": "classifier"},
    "yesno_obzv": {"text": "да", "type": "classifier"},
    "weekday": {"text": "в пятницу", "type": "classifier"},
    "end_repeat": {"text": "все, до свидания", "type": "classifier"},
    "operatorcheck": {"text": "соедините с оператором", "type": "classifier"},
    "see_grade": {"text": "пять", "type": "classifier"},
    "other_payment": {"text": "мобилизационные выплаты", "type": "classifier"},
    "mobilizac": {"text": "мобилизационные", "type": "classifier"},
    "matcap_direct": {"text": "на улучшение жилищных условий", "type": "classifier"},
    "new_minsoc_subsidii": {"text": "субсидия на жкх", "type": "classifier"},
    "new_trud": {"text": "потерял работу", "type": "classifier"},
    "rating_med": {"text": "отлично", "type": "classifier"},
    "zags": {"text": "загс центрального района", "type": "classifier"},
    "zags_handler": {"text": "загс на нижегородской", "type": "classifier"},
    "filial": {"text": "в мошковский филиал", "type": "classifier"},
    "offices_mfc_see": {"text": "родники", "type": "classifier"},
    "service_mfc_see": {"text": "хочу получить водительское удостоверение", "type": "classifier"},
    "nsk_fillials": {"text": "в филиал на родниках", "type": "classifier"},
}

CUSTOM_MODEL_MATCHES = [
    ("record_date", "на завтра"),
    ("record_time", "в пять часов вечера"),
    ("fio", "Иванов Иван Иванович"),
    ("phone", "8 913 456 78 90"),
    ("date", "на завтра"),
    ("time", "в пять часов вечера"),
    ("uuid", "550e8400-e29b-41d4-a716-446655440000"),
    ("number", "пятьсот двадцать три"),
    ("filial", "в филиал на родниках"),
    ("zags", "загс центрального района"),
    ("yesno", "да"),
    ("weekday", "в пятницу"),
    ("grade", "пять"),
    ("obzv", "да"),
    ("operator", "соедините с оператором"),
    ("repeat", "все, до свидания"),
    ("matcap", "на улучшение жилищных условий"),
    ("mobilizac", "мобилизационные"),
    ("trud", "потерял работу"),
    ("subsidii", "субсидия на жкх"),
    ("rating", "отлично"),
    ("payment", "мобилизационные выплаты"),
    ("service_mfc", "хочу получить паспорт"),
    ("offices_mfc", "родники"),
    ("паспорт", "хочу получить загранпаспорт"),
    ("ЕПГУ", "проверить статус заявления на госуслугах"),
    ("лаврентьев", "как получить лаврентьевскую карту"),
    ("возраст", "мне 25 лет"),
    ("загран", "хочу сделать загранпаспорт"),
    ("вернуться", "вернуться в главное меню"),
    ("не знаю код", "не знаю код подразделения"),
    ("оценк", "оцениваю обслуживание на пять"),
    ("несогласи", "я не согласен с этим"),
    ("начало диалог", "здравствуйте"),
    ("пособие", "какое пособие по рождению"),
    ("УК", "проблема с управляющей компанией"),
]

FALLBACK_TEXT = "Здравствуйте, хочу записаться на приём в МФЦ для получения паспорта, запись на завтра"


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=15):
        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):
        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["models"]) + len(data["handlers"]) + len(data["pretrained"])
        print(f"\n  models: {len(data['models'])}, "
              f"handlers: {len(data['handlers'])}, "
              f"pretrained: {len(data['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"] 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 _text_for_model(self, model_id):
        a = self._find_assert(model_id)
        if a and "text" in a:
            return a["text"]
        for substr, text in CUSTOM_MODEL_MATCHES:
            if substr in model_id:
                return text
        return FALLBACK_TEXT

    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 test_01_list_models(self):
        """Получение списка моделей"""
        data = self.models
        self.assertIn("models", data)
        self.assertIn("handlers", data)
        self.assertIn("pretrained", data)

    def _fmt_result(self, model_id, result):
        print(f"  >> {model_id}")
        if isinstance(result, dict) and "classes" in result:
            for c in result["classes"][:5]:
                print(f"      class={c.get('class','?')} conf={c.get('confidence',0):.3f}")
            if len(result["classes"]) > 5:
                print(f"      ... ({len(result['classes'])} classes)")
        else:
            for key, entities in result.items():
                if isinstance(entities, list):
                    for e in entities[:3]:
                        text = e.get("text", "?")
                        calc = e.get("calculated", "?")
                        conf = e.get("confidence", 0)
                        print(f"      {key}: text='{text}' calc='{calc}' conf={conf:.3f}")
                    if len(entities) > 3:
                        print(f"      ... ({len(entities)} entities)")

    def test_02_pretrained_models(self):
        """Проверка pretrained-моделей"""
        errors = []
        for model_id, td in PRETRAINED_TESTS.items():
            if model_id not in self.models.get("pretrained", []):
                errors.append(f"  {model_id}: нет в списке")
                continue
            try:
                result = self.api.test_entity(model_id, td["text"])
                if self._is_empty(result):
                    errors.append(f"  {model_id}: пустой результат")
                else:
                    self._check_assert(model_id, result)
                    self._fmt_result(model_id, result)
            except AssertionError:
                raise
            except Exception as e:
                errors.append(f"  {model_id}: {e}")
        if errors:
            self.fail("Ошибки pretrained:\n" + "\n".join(errors))

    def test_03_handler_models(self):
        """Проверка handler-моделей"""
        errors = []
        handlers = self.models.get("handlers", [])
        if not handlers:
            self.skipTest("Нет handler моделей")
        for model_id in handlers:
            text = FALLBACK_TEXT
            for key, td in HANDLER_TESTS.items():
                if key in model_id:
                    text = td["text"]
                    break
            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):
                    print(f"  [!] {model_id}: пустой результат")
                else:
                    self._check_assert(model_id, result)
                    self._fmt_result(model_id, result)
            except AssertionError:
                raise
            except Exception as e:
                errors.append(f"  {model_id}: {e}")
        if errors:
            self.fail("Ошибки handler:\n" + "\n".join(errors))

    def test_04_custom_models(self):
        """Проверка кастомных моделей"""
        errors = []
        models = self.models.get("models", [])
        if not models:
            self.skipTest("Нет кастомных моделей")
        for model_id in models:
            text = self._text_for_model(model_id)
            try:
                result = self.api.test_entity(model_id, text)
                if self._is_empty(result):
                    print(f"  [!] {model_id}")
                else:
                    self._check_assert(model_id, result)
                    self._fmt_result(model_id, result)
            except AssertionError:
                raise
            except Exception as e:
                errors.append(f"  {model_id}: {e}")
        if errors:
            self.fail("Ошибки кастомных моделей:\n" + "\n".join(errors))

    def test_05_handler_pretrained_coverage(self):
        """Проверка handler моделей phone/uuid/number/fio"""
        errors = []
        handlers = self.models.get("handlers", [])
        for htype, td in [("phone", PRETRAINED_TESTS["phone"]),
                          ("uuid", PRETRAINED_TESTS["uuid"]),
                          ("number", PRETRAINED_TESTS["number"]),
                          ("fio", PRETRAINED_TESTS["fio"])]:
            matching = [h for h in handlers if htype in h.lower()]
            if not matching:
                continue
            for mid in matching[:2]:
                try:
                    result = self.api.test_entity(mid, td["text"])
                    if self._is_empty(result):
                        print(f"  [!] {mid} ({td['hint']})")
                    else:
                        self._check_assert(mid, result)
                        self._fmt_result(mid, result)
                except AssertionError:
                    raise
                except Exception as e:
                    errors.append(f"  {mid}: {e}")
        if errors:
            self.fail("Ошибки handler coverage:\n" + "\n".join(errors))


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 файл с ассертами")
    args = parser.parse_args()
    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()