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

Использование:
    python3 test_smc_models.py
    python3 test_smc_models.py --host 10.101.60.27
    python3 test_smc_models.py --host 10.101.60.27 --port 6181
    python3 test_smc_models.py --asserts smc_asserts.json

Файл ассертов (JSON):
    {
      "asserts": [
        {
          "model": "МФЦ НСК",
          "text": "записаться на приём",   // опционально, свой текст
          "checks": {
            "min_found": 1,               // мин. количество классов
            "text_contains": "Запись",     // подстрока в top классе
            "confidence_min": 0.5          // мин. confidence топ-класса
          }
        }
      ]
    }

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

import argparse
import json
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.27"
DEFAULT_PORT = 6181
BASE_PATH = "/smc"

TEXT_RULES = [
    ("филиал", "родники"),
    ("offices", "родники"),
    ("мфц", "записаться в мфц"),
    ("mfc", "записаться в мфц"),
    ("ekc", "записаться в екц"),
    ("екц", "записаться в екц"),
    ("minsoc", "какие выплаты"),
    ("минсоц", "какие выплаты"),
    ("техническая", "подключите оператора"),
    ("поддержка", "подключите оператора"),
    ("тест", "тестовый запрос"),
    ("test", "test query"),
]


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 SmcAPI:
    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 classify(self, model_id, text):
        encoded = quote(model_id, safe="")
        url = f"{self.base}/classify/{encoded}?text={quote(text)}"
        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 TestSmcModels(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.api = SmcAPI(cls.args.host, cls.args.port)
        cls.asserts = load_asserts(cls.args.asserts)
        data = cls.api.get_models()
        TestSmcModels.models = data
        models = data.get("models", [])
        handlers = data.get("handlers", [])
        print(f"\n  models: {len(models)}, handlers: {len(handlers)}")
        if cls.asserts:
            print(f"  asserts loaded: {len(cls.asserts)}")
        if len(models) + len(handlers) == 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", {})
        classes = result.get("classes", [])
        if "min_found" in checks:
            self.assertGreaterEqual(len(classes), checks["min_found"],
                f"{model_id}: expected >= {checks['min_found']} classes, got {len(classes)}")
        if classes and "confidence_min" in checks:
            self.assertGreaterEqual(classes[0].get("confidence", 0), checks["confidence_min"],
                f"{model_id}: top confidence < {checks['confidence_min']}")

    def test_01_list_models(self):
        """Получение списка моделей"""
        data = self.models
        self.assertIn("models", data)
        self.assertIn("handlers", data)
        print(f"\n  models: {len(data['models'])}")
        print(f"  handlers: {len(data['handlers'])}")

    def _fmt_result(self, result):
        classes = result.get("classes", [])
        if classes:
            top = classes[0]
            return f"{top.get('class','?')} ({top.get('confidence',0):.3f})"
        return "∅"

    def _run_section(self, section_name, model_ids):
        results = []
        for mid in model_ids:
            if not mid:
                continue
            text = pick_text(mid)
            a = self._find_assert(mid)
            if a and "text" in a:
                text = a["text"]
            try:
                result = self.api.classify(mid, text)
                classes = result.get("classes", [])
                if not classes:
                    results.append(("empty", mid, None, None))
                else:
                    self._check_assert(mid, result)
                    results.append(("ok", mid, result, None))
            except AssertionError:
                results.append(("assert_fail", mid, None, None))
                raise
            except Exception as e:
                results.append(("error", mid, None, str(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"]

        print(f"\n── {section_name} ──────────────────────────────────────────────")
        for status, mid, result, extra in results:
            if status == "ok" and self.args.verbose:
                print(f"  ✓ {mid}  → {self._fmt_result(result)}")
            elif status == "ok":
                print(f"  ✓ {mid}")
            elif status == "empty":
                print(f"  ✗ {mid}  → empty")
            elif status == "assert_fail":
                print(f"  ✗ {mid}  → assert fail")
            elif status == "error":
                print(f"  ⚠ {mid}  → {extra}")
        print(f"  ─── {passed} passed, {empty} empty", end="")
        if errors:
            print(f", {len(errors)} errors", end="")
        print()

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

    def test_02_models(self):
        """Классификация через модели"""
        models = self.models.get("models", [])
        if not models:
            self.skipTest("Нет моделей")
        self._run_section("MODELS", models)

    def test_03_handlers(self):
        """Классификация через handler модели"""
        handlers = self.models.get("handlers", [])
        if not handlers:
            self.skipTest("Нет handler моделей")
        self._run_section("HANDLERS", handlers)


def main():
    parser = argparse.ArgumentParser(description="SMC 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="Показать результат каждой модели")
    args = parser.parse_args()
    TestSmcModels.args = args

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


if __name__ == "__main__":
    main()