#!/usr/bin/env python3
"""
Отправка потока кадров (одно изображение многократно) в SDK ФПГ (GigaDoc)
для получения полных медицинских показателей.

По документации: ~300 кадров JPEG при ~30 кадрах/сек (~10 сек видео).
"""

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

import requests
import websockets

# ─── Конфигурация ─────────────────────────────────────────────
OAUTH_URL = "https://oauth.sbermed.ai/token"
CLIENT_ID = "client_fpg_front"
CLIENT_SECRET = os.environ.get("FPG_CLIENT_SECRET", "")

REST_HOST = "gigadoc-sdk.sbermed.ai"
REST_BASE = f"https://{REST_HOST}"

WS_HOST = "gigadoc-sdk.sbermed.ai"
WS_PORT = 8443
WS_PATH_TEMPLATE = "/api/v2/ws/predict/{uid}/image"

DEVICE_ID = "web_prod"
FORM_FACTOR = "web_prod"

LOG_FILE = "fpg_stream.log"


# ─── Логирование ──────────────────────────────────────────────
def setup_logging(log_file: str = LOG_FILE):
    logger = logging.getLogger("fpg_stream")
    logger.setLevel(logging.DEBUG)

    formatter = logging.Formatter(
        "%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    fh = logging.FileHandler(log_file, encoding="utf-8")
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(formatter)
    logger.addHandler(fh)

    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)
    ch.setFormatter(formatter)
    logger.addHandler(ch)

    return logger


logger = setup_logging()


# ─── Этап 1: Авторизация ──────────────────────────────────────
def get_token():
    logger.info("=" * 60)
    logger.info("ЭТАП 1: Авторизация")
    logger.info("=" * 60)
    logger.info(f"POST {OAUTH_URL}")

    resp = requests.post(
        OAUTH_URL,
        data={
            "grant_type": "client_credentials",
            "client_id": CLIENT_ID,
            "client_secret": CLIENT_SECRET,
        },
        timeout=15,
    )
    logger.info(f"HTTP {resp.status_code}")

    if resp.status_code != 200:
        raise RuntimeError(f"Ошибка авторизации: {resp.status_code} {resp.text}")

    data = resp.json()
    token = data["access_token"]
    logger.info(f"✅ Токен получен (действует {data.get('expires_in', '?')} сек)")
    return token


# ─── Этап 2: Создание сессии ──────────────────────────────────
def create_session(token: str):
    logger.info("=" * 60)
    logger.info("ЭТАП 2: Создание сессии")
    logger.info("=" * 60)
    url = f"{REST_BASE}/api/v1/predict/new"
    logger.info(f"POST {url}")
    logger.info(f"  device_id={DEVICE_ID}")
    logger.info(f"  form_factor={FORM_FACTOR}")

    resp = requests.post(
        url,
        params={"device_id": DEVICE_ID, "form_factor": FORM_FACTOR},
        headers={"Authorization": f"Bearer {token}"},
        timeout=15,
    )
    logger.info(f"HTTP {resp.status_code}")

    if resp.status_code != 200:
        raise RuntimeError(f"Ошибка создания сессии: {resp.status_code} {resp.text}")

    uid = resp.json()["uid"]
    logger.info(f"✅ Сессия создана: {uid}")
    return uid


# ─── Этап 3-5: WebSocket + отправка потока + результат ────────
async def ws_stream(token: str, uid: str, image_path: str, frames: int, fps: float):
    logger.info("=" * 60)
    logger.info("ЭТАП 3: WebSocket подключение")
    logger.info("=" * 60)

    uri = f"wss://{WS_HOST}:{WS_PORT}{WS_PATH_TEMPLATE.format(uid=uid)}"
    logger.info(f"URL: {uri}")

    # Читаем изображение (или директорию с кадрами)
    import glob

    if os.path.isdir(image_path):
        # Директория с кадрами
        patterns = ["*.jpg", "*.jpeg", "*.png"]
        frame_files = []
        for p in patterns:
            frame_files.extend(glob.glob(os.path.join(image_path, p)))
        frame_files.sort()
        logger.info(f"Директория с кадрами: {image_path} ({len(frame_files)} файлов)")
        if not frame_files:
            logger.error("Нет кадров в директории")
            return None
        # Ограничиваем количество кадров
        frame_files = frame_files[:frames]
        logger.info(f"Будет отправлено кадров: {len(frame_files)}")
    else:
        # Один файл — повторяем его frames раз
        frame_files = [image_path] * frames
        logger.info(f"Изображение: {image_path} ({os.path.getsize(image_path)} bytes), повторяем {frames} раз")

    logger.info(f"Частота: {fps} fps, интервал: {1000/fps:.0f} мс")

    interval = 1.0 / fps

    try:
        async with websockets.connect(
            uri,
            open_timeout=10,
            subprotocols=[f"bearer.{token}"],
            origin="https://gigadoc.sbermed.ai",
            extra_headers={
                "Cache-Control": "no-cache",
                "Pragma": "no-cache",
            },
            compression="deflate",
            ping_interval=20,
            ping_timeout=20,
            max_size=None,
        ) as ws:
            logger.info(f"✅ WebSocket подключен: {uri}")
            logger.info(f"  Subprotocol: {ws.subprotocol}")

            logger.info("=" * 60)
            logger.info("ЭТАП 4: Отправка потока кадров")
            logger.info("=" * 60)

            result = None
            last_progress = -1
            t_start = time.time()

            for i, frame_file in enumerate(frame_files):
                frame_start = time.time()

                # Читаем кадр
                with open(frame_file, "rb") as f:
                    img_data = f.read()

                # Отправляем кадр
                await ws.send(img_data)

                # Читаем ответ (не блокируемся долго)
                try:
                    resp = await asyncio.wait_for(ws.recv(), timeout=2.0)
                    try:
                        data = json.loads(resp)
                        progress = data.get("progress", 0.0)
                        face_ok = data.get("face_ok", False)

                        # Логируем прогресс при изменении
                        pct = int(progress * 100)
                        if pct != last_progress:
                            logger.info(f"  [{i+1}/{frames}] progress={progress:.2f} ({pct}%) face_ok={face_ok}")
                            last_progress = pct

                        # Финальный результат
                        if progress >= 1.0:
                            logger.info("=" * 60)
                            logger.info("ЭТАП 5: ФИНАЛЬНЫЙ РЕЗУЛЬТАТ")
                            logger.info("=" * 60)
                            logger.info(json.dumps(data, indent=2, ensure_ascii=False))
                            result = data
                            break
                    except json.JSONDecodeError:
                        logger.debug(f"  Ответ не JSON: {resp[:200]}")
                except asyncio.TimeoutError:
                    pass  # нет ответа на этот кадр — нормально

                # Соблюдаем частоту кадров
                elapsed = time.time() - frame_start
                if elapsed < interval:
                    await asyncio.sleep(interval - elapsed)

                # Прогресс-индикатор каждые 50 кадров
                if (i + 1) % 50 == 0:
                    elapsed_total = time.time() - t_start
                    logger.info(f"  ... отправлено {i+1}/{frames} кадров за {elapsed_total:.1f} сек")

            if result is None:
                logger.warning(f"Финальный результат не получен после {frames} кадров")
                logger.info(f"  Отправлено кадров: {frames}, время: {time.time()-t_start:.1f} сек")
            return result

    except Exception as e:
        logger.error(f"❌ Ошибка WebSocket: {type(e).__name__}: {e}")
        return None


# ─── Этап 6: Удаление сессии ──────────────────────────────────
def delete_session(token: str, uid: str):
    logger.info("=" * 60)
    logger.info("ЭТАП 6: Удаление сессии")
    logger.info("=" * 60)
    url = f"{REST_BASE}/api/v1/predict/{uid}"
    logger.info(f"DELETE {url}")

    resp = requests.delete(
        url,
        headers={"Authorization": f"Bearer {token}"},
        timeout=15,
    )
    logger.info(f"HTTP {resp.status_code}")

    if resp.status_code == 200:
        logger.info(f"✅ Сессия {uid} удалена")
    else:
        logger.warning(f"⚠️ Сессия {uid} не удалена: {resp.status_code} {resp.text}")


# ─── Main ─────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(description="Отправка потока кадров в SDK ФПГ (GigaDoc)")
    parser.add_argument("--source", required=True, help="Путь к JPEG-изображению или директории с кадрами")
    parser.add_argument("--frames", type=int, default=300, help="Количество кадров (default: 300)")
    parser.add_argument("--fps", type=float, default=30.0, help="Частота кадров (default: 30)")
    parser.add_argument("--log-file", default=LOG_FILE, help=f"Файл лога (default: {LOG_FILE})")
    args = parser.parse_args()

    # Настройка логирования
    global logger
    logger = setup_logging(args.log_file)

    logger.info("=" * 60)
    logger.info("🏥 Отправка потока кадров в SDK ФПГ (GigaDoc)")
    logger.info(f"Дата: {datetime.now().isoformat()}")
    logger.info(f"Источник: {args.source}")
    logger.info(f"Кадры: {args.frames} @ {args.fps} fps")
    logger.info("=" * 60)

    if not os.path.exists(args.source):
        logger.error(f"Файл/директория не найден: {args.source}")
        return 1

    # 1. Авторизация
    try:
        token = get_token()
    except Exception as e:
        logger.error(f"Тест прерван: {e}")
        return 1

    # 2. Создание сессии
    try:
        uid = create_session(token)
    except Exception as e:
        logger.error(f"Тест прерван: {e}")
        return 1

    # 3-5. WebSocket + поток
    result = asyncio.run(ws_stream(token, uid, args.source, args.frames, args.fps))

    # 6. Удаление сессии
    delete_session(token, uid)

    # Итог
    logger.info("=" * 60)
    if result:
        logger.info("✅ Тест завершён успешно — полный результат получен")
    else:
        logger.info("⚠️ Результат не получен (сервер не завершил диагностику)")
    logger.info(f"Лог сохранён в: {args.log_file}")
    logger.info("=" * 60)

    return 0 if result else 2


if __name__ == "__main__":
    sys.exit(main())