#!/usr/bin/env python3

import argparse
import asyncio
import json
import logging
import ssl
import sys
from pathlib import Path
from typing import Optional

import aiohttp
import websockets
from websockets.exceptions import (
    ConnectionClosed,
    InvalidMessage,
    InvalidStatus,
)


LOG = logging.getLogger("gigadoc-client")


# ============================================================
# CLI
# ============================================================

def parse_args():
    parser = argparse.ArgumentParser(
        description=(
            "GigaDoc SDK client: "
            "Keycloak -> REST -> WebSocket -> DELETE"
        )
    )

    # --------------------------------------------------------
    # Keycloak
    # --------------------------------------------------------

    parser.add_argument(
        "--keycloak-url",
        required=True,
        help=(
            "Keycloak token URL, например: "
            "https://keycloak.example.com/realms/REALM/"
            "protocol/openid-connect/token"
        ),
    )

    parser.add_argument(
        "--client-id",
        required=True,
        help="Keycloak client_id",
    )

    parser.add_argument(
        "--client-secret",
        required=True,
        help="Keycloak client_secret",
    )

    # --------------------------------------------------------
    # SDK
    # --------------------------------------------------------

    parser.add_argument(
        "--sdk-host",
        required=True,
        help="SDK hostname, например gigadoc-sdk.sbermed.ai",
    )

    parser.add_argument(
        "--rest-prefix",
        default="/api/v1",
        help="REST API prefix (default: /api/v1)",
    )

    parser.add_argument(
        "--ws-prefix",
        default="/api/v2/ws/predict",
        help="WebSocket API prefix (default: /api/v2/ws/predict)",
    )

    parser.add_argument(
        "--ws-port",
        type=int,
        action="append",
        default=None,
        help=(
            "WebSocket port. Можно указать несколько раз. "
            "По умолчанию: 8443,443"
        ),
    )

    parser.add_argument(
        "--origin",
        default="https://gigadoc.sbermed.ai",
        help=(
            "Origin для WebSocket "
            "(default: https://gigadoc.sbermed.ai)"
        ),
    )

    # --------------------------------------------------------
    # Predict
    # --------------------------------------------------------

    parser.add_argument(
        "--device-id",
        required=True,
        help="device_id",
    )

    parser.add_argument(
        "--form-factor",
        required=True,
        help="form_factor",
    )

    # --------------------------------------------------------
    # Image
    # --------------------------------------------------------

    parser.add_argument(
        "--image",
        required=True,
        help="Путь к JPEG изображению",
    )

    # --------------------------------------------------------
    # Timeouts
    # --------------------------------------------------------

    parser.add_argument(
        "--http-timeout",
        type=float,
        default=30.0,
        help="HTTP timeout в секундах (default: 30)",
    )

    parser.add_argument(
        "--ws-open-timeout",
        type=float,
        default=30.0,
        help="WebSocket handshake timeout (default: 30)",
    )

    parser.add_argument(
        "--ws-message-timeout",
        type=float,
        default=60.0,
        help=(
            "Timeout ожидания следующего WS сообщения "
            "(default: 60)"
        ),
    )

    # --------------------------------------------------------
    # Options
    # --------------------------------------------------------

    parser.add_argument(
        "--insecure",
        action="store_true",
        help="Отключить проверку TLS сертификата",
    )

    parser.add_argument(
        "--keep-alive",
        action="store_true",
        help=(
            "Не закрывать WS сразу после progress=1.0; "
            "продолжать читать сообщения"
        ),
    )

    parser.add_argument(
        "--no-delete",
        action="store_true",
        help="Не удалять predict session после завершения",
    )

    parser.add_argument(
        "--verbose",
        action="store_true",
        help="DEBUG logging",
    )

    return parser.parse_args()


# ============================================================
# Logging
# ============================================================

def setup_logging(verbose: bool):
    logging.basicConfig(
        level=logging.DEBUG if verbose else logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
    )


# ============================================================
# Helpers
# ============================================================

def normalize_prefix(prefix: str) -> str:
    if not prefix.startswith("/"):
        prefix = "/" + prefix

    return prefix.rstrip("/")


def build_ssl_context(
    insecure: bool,
) -> Optional[ssl.SSLContext]:
    """
    Для обычного wss:// возвращаем None.

    В websockets 17.1 параметр ssl=None нельзя явно
    передавать вместе с wss://.

    Поэтому SSL context добавляется только при --insecure.
    """

    if not insecure:
        return None

    context = ssl.create_default_context()

    context.check_hostname = False
    context.verify_mode = ssl.CERT_NONE

    return context


# ============================================================
# Keycloak
# ============================================================

async def get_access_token(
    session: aiohttp.ClientSession,
    args,
) -> str:

    LOG.info("")
    LOG.info("=" * 70)
    LOG.info("KEYCLOAK")
    LOG.info("=" * 70)

    LOG.info(
        "POST %s",
        args.keycloak_url,
    )

    data = {
        "grant_type": "client_credentials",
        "client_id": args.client_id,
        "client_secret": args.client_secret,
    }

    try:
        async with session.post(
            args.keycloak_url,
            data=data,
        ) as response:

            body = await response.text()

            LOG.info(
                "Keycloak -> HTTP %s",
                response.status,
            )

            if args.verbose:
                LOG.debug(
                    "Keycloak response body: %s",
                    body,
                )

            if response.status != 200:
                raise RuntimeError(
                    "Keycloak returned HTTP "
                    f"{response.status}: {body}"
                )

            try:
                payload = json.loads(body)
            except json.JSONDecodeError as exc:
                raise RuntimeError(
                    "Keycloak returned invalid JSON"
                ) from exc

            access_token = payload.get(
                "access_token"
            )

            if not access_token:
                raise RuntimeError(
                    "Keycloak response does not contain "
                    "access_token"
                )

            LOG.info(
                "Access token получен"
            )

            return access_token

    except Exception:
        LOG.exception(
            "Ошибка получения access token"
        )
        raise


# ============================================================
# Create session
# ============================================================

async def create_predict_session(
    session: aiohttp.ClientSession,
    args,
    access_token: str,
) -> str:

    rest_prefix = normalize_prefix(
        args.rest_prefix
    )

    url = (
        f"https://{args.sdk_host}"
        f"{rest_prefix}/predict/new"
    )

    params = {
        "device_id": args.device_id,
        "form_factor": args.form_factor,
    }

    headers = {
        "Authorization": f"Bearer {access_token}",
        "Accept": "application/json",
    }

    LOG.info("")
    LOG.info("=" * 70)
    LOG.info("CREATE PREDICT SESSION")
    LOG.info("=" * 70)

    LOG.info(
        "POST %s",
        url,
    )

    LOG.info(
        "device_id=%s",
        args.device_id,
    )

    LOG.info(
        "form_factor=%s",
        args.form_factor,
    )

    try:
        async with session.post(
            url,
            params=params,
            headers=headers,
        ) as response:

            body = await response.text()

            LOG.info(
                "POST /predict/new -> HTTP %s",
                response.status,
            )

            if args.verbose:
                LOG.debug(
                    "Response body: %s",
                    body,
                )

            if not (
                200 <= response.status < 300
            ):
                raise RuntimeError(
                    f"POST {url} returned HTTP "
                    f"{response.status}: {body}"
                )

            try:
                payload = json.loads(body)
            except json.JSONDecodeError as exc:
                raise RuntimeError(
                    "Invalid JSON from predict/new: "
                    f"{body}"
                ) from exc

            uid = payload.get("uid")

            if not uid:
                raise RuntimeError(
                    "Response does not contain uid: "
                    f"{payload}"
                )

            LOG.info(
                "UID = %s",
                uid,
            )

            return uid

    except Exception:
        LOG.exception(
            "Ошибка создания predict session"
        )
        raise


# ============================================================
# WebSocket
# ============================================================

async def websocket_predict(
    args,
    access_token: str,
    uid: str,
) -> bool:

    ws_prefix = normalize_prefix(
        args.ws_prefix
    )

    # Если --ws-port не задан:
    # сначала 8443, затем 443.
    ports = args.ws_port

    if not ports:
        ports = [
            8443,
            443,
        ]

    image_path = Path(
        args.image
    )

    if not image_path.is_file():
        raise FileNotFoundError(
            f"Image not found: {image_path}"
        )

    image_data = image_path.read_bytes()

    if not image_data:
        raise RuntimeError(
            f"Image is empty: {image_path}"
        )

    LOG.info("")
    LOG.info(
        "JPEG: %s",
        image_path,
    )

    LOG.info(
        "JPEG size: %d bytes",
        len(image_data),
    )

    # --------------------------------------------------------
    # ВАЖНО:
    #
    # Фронт подключается примерно так:
    #
    # Sec-WebSocket-Protocol:
    #     bearer.<ACCESS_TOKEN>
    #
    # Поэтому делаем абсолютно то же самое.
    # --------------------------------------------------------

    bearer_protocol = (
        f"bearer.{access_token}"
    )

    ssl_context = build_ssl_context(
        args.insecure
    )

    for port in ports:

        ws_url = (
            f"wss://{args.sdk_host}:{port}"
            f"{ws_prefix}/{uid}/image"
        )

        LOG.info("")
        LOG.info("=" * 70)
        LOG.info(
            "WEBSOCKET :%d",
            port,
        )
        LOG.info("=" * 70)

        LOG.info(
            "URL: %s",
            ws_url,
        )

        LOG.debug(
            "Origin: %s",
            args.origin,
        )

        LOG.debug(
            "Subprotocol: bearer.<TOKEN>",
        )

        # ----------------------------------------------------
        # Формируем kwargs отдельно.
        #
        # Нельзя передавать ssl=None в websockets 17.1
        # для wss://.
        # ----------------------------------------------------

        connect_kwargs = {
            "subprotocols": [
                bearer_protocol
            ],

            "additional_headers": {
                "Origin": args.origin,
                "Cache-Control": "no-cache",
                "Pragma": "no-cache",
            },

            "compression": "deflate",

            "open_timeout": (
                args.ws_open_timeout
            ),

            "ping_interval": 20,

            "ping_timeout": 20,

            "max_size": None,
        }

        # Только при --insecure.
        if ssl_context is not None:
            connect_kwargs["ssl"] = (
                ssl_context
            )

        try:

            LOG.debug(
                "Открываем WebSocket..."
            )

            async with websockets.connect(
                ws_url,
                **connect_kwargs,
            ) as ws:

                LOG.info("")
                LOG.info(
                    "✅ WEBSOCKET :%d CONNECTED",
                    port,
                )

                LOG.info(
                    "Negotiated subprotocol: %s",
                    ws.subprotocol,
                )

                # ------------------------------------------------
                # Проверяем, что сервер согласовал bearer protocol.
                # ------------------------------------------------

                if ws.subprotocol != bearer_protocol:

                    LOG.warning(
                        "⚠️ Server did not negotiate "
                        "expected bearer subprotocol"
                    )

                    LOG.warning(
                        "Expected: bearer.<TOKEN>"
                    )

                    LOG.warning(
                        "Received: %s",
                        ws.subprotocol,
                    )

                else:
                    LOG.info(
                        "✅ Bearer subprotocol accepted"
                    )

                # ------------------------------------------------
                # Отправляем JPEG как binary frame
                # ------------------------------------------------

                LOG.info("")
                LOG.info(
                    "Отправка JPEG..."
                )

                await ws.send(
                    image_data
                )

                LOG.info(
                    "✅ JPEG отправлен: %d bytes",
                    len(image_data),
                )

                # ------------------------------------------------
                # Читаем ответы
                # ------------------------------------------------

                while True:

                    try:
                        message = await asyncio.wait_for(
                            ws.recv(),
                            timeout=(
                                args.ws_message_timeout
                            ),
                        )

                    except asyncio.TimeoutError:

                        LOG.error(
                            "⏱ Таймаут ожидания "
                            "сообщения от WebSocket"
                        )

                        return False

                    except ConnectionClosed as exc:

                        LOG.warning(
                            "WebSocket закрыт сервером: %s",
                            exc,
                        )

                        return False

                    if message is None:

                        LOG.warning(
                            "WebSocket returned None"
                        )

                        return False

                    # ------------------------------------------------
                    # Binary response
                    # ------------------------------------------------

                    if isinstance(
                        message,
                        bytes,
                    ):

                        LOG.info(
                            "WS <- binary: %d bytes",
                            len(message),
                        )

                        continue

                    # ------------------------------------------------
                    # JSON response
                    # ------------------------------------------------

                    LOG.debug(
                        "WS <- %s",
                        message,
                    )

                    try:
                        payload = json.loads(
                            message
                        )

                    except json.JSONDecodeError:

                        LOG.warning(
                            "Получено не-JSON "
                            "сообщение:"
                        )

                        LOG.warning(
                            "%s",
                            message,
                        )

                        continue

                    LOG.info(
                        "WS result: %s",
                        json.dumps(
                            payload,
                            ensure_ascii=False,
                        ),
                    )

                    # ------------------------------------------------
                    # Progress
                    #
                    # Согласно API:
                    #
                    # 0.0 ... 1.0
                    #
                    # 1.0 = завершено
                    # ------------------------------------------------

                    progress = payload.get(
                        "progress"
                    )

                    if progress is None:
                        continue

                    try:

                        progress_float = float(
                            progress
                        )

                    except (
                        TypeError,
                        ValueError,
                    ):

                        LOG.warning(
                            "Некорректный progress: %r",
                            progress,
                        )

                        continue

                    LOG.info(
                        "Progress: %.1f%%",
                        progress_float * 100,
                    )

                    if progress_float >= 1.0:

                        LOG.info("")
                        LOG.info(
                            "✅ PREDICT ЗАВЕРШЁН"
                        )

                        # ------------------------------------------------
                        # Если keep-alive не задан — сразу выходим.
                        # ------------------------------------------------

                        if not args.keep_alive:
                            return True

                        LOG.info(
                            "--keep-alive задан, "
                            "оставляем WebSocket открытым"
                        )

                        while True:

                            try:

                                message = (
                                    await asyncio.wait_for(
                                        ws.recv(),
                                        timeout=(
                                            args.ws_message_timeout
                                        ),
                                    )
                                )

                                LOG.info(
                                    "WS <- %s",
                                    message,
                                )

                            except asyncio.TimeoutError:

                                LOG.info(
                                    "Keep-alive timeout"
                                )

                                return True

                            except ConnectionClosed as exc:

                                LOG.info(
                                    "Server closed WS: %s",
                                    exc,
                                )

                                return True

        # --------------------------------------------------------
        # HTTP status ошибки handshake
        # --------------------------------------------------------

        except InvalidStatus as exc:

            LOG.error("")
            LOG.error(
                "❌ WEBSOCKET :%d FAILED",
                port,
            )

            LOG.error(
                "Exception type: %s",
                type(exc).__name__,
            )

            LOG.error(
                "Exception: %s",
                exc,
            )

            continue

        # --------------------------------------------------------
        # Сервер закрыл соединение без HTTP response.
        # --------------------------------------------------------

        except InvalidMessage as exc:

            LOG.error("")
            LOG.error(
                "❌ WEBSOCKET :%d FAILED",
                port,
            )

            LOG.error(
                "Exception type: %s",
                type(exc).__name__,
            )

            LOG.error(
                "Exception: %s",
                exc,
            )

            continue

        # --------------------------------------------------------
        # WS закрылся после подключения
        # --------------------------------------------------------

        except ConnectionClosed as exc:

            LOG.error("")
            LOG.error(
                "❌ WEBSOCKET :%d CLOSED",
                port,
            )

            LOG.error(
                "Exception: %s",
                exc,
            )

            continue

        # --------------------------------------------------------
        # Остальные ошибки
        # --------------------------------------------------------

        except Exception as exc:

            LOG.error("")
            LOG.error(
                "❌ WEBSOCKET :%d FAILED",
                port,
            )

            LOG.error(
                "Exception type: %s",
                type(exc).__name__,
            )

            LOG.error(
                "Exception: %s",
                exc,
            )

            continue

    return False


# ============================================================
# DELETE
# ============================================================

async def delete_predict_session(
    session: aiohttp.ClientSession,
    args,
    access_token: str,
    uid: str,
):

    rest_prefix = normalize_prefix(
        args.rest_prefix
    )

    url = (
        f"https://{args.sdk_host}"
        f"{rest_prefix}/predict/{uid}"
    )

    headers = {
        "Authorization": f"Bearer {access_token}",
    }

    LOG.info("")
    LOG.info("=" * 70)
    LOG.info("DELETE PREDICT SESSION")
    LOG.info("=" * 70)

    LOG.info(
        "DELETE %s",
        url,
    )

    try:

        async with session.delete(
            url,
            headers=headers,
        ) as response:

            body = await response.text()

            LOG.info(
                "DELETE /predict/%s -> HTTP %s",
                uid,
                response.status,
            )

            if args.verbose and body:
                LOG.debug(
                    "Delete response: %s",
                    body,
                )

            if not (
                200 <= response.status < 300
            ):

                LOG.warning(
                    "⚠️ DELETE failed: HTTP %s",
                    response.status,
                )

                if body:
                    LOG.warning(
                        "Response: %s",
                        body,
                    )

            else:

                LOG.info(
                    "✅ Session удалена"
                )

    except Exception:

        LOG.exception(
            "Ошибка удаления predict session"
        )


# ============================================================
# Main
# ============================================================

async def main():

    args = parse_args()

    setup_logging(
        args.verbose
    )

    LOG.info("")
    LOG.info("=" * 70)
    LOG.info("GIGADOC SDK PYTHON CLIENT")
    LOG.info("=" * 70)

    LOG.info(
        "SDK host: %s",
        args.sdk_host,
    )

    LOG.info(
        "REST prefix: %s",
        args.rest_prefix,
    )

    LOG.info(
        "WS prefix: %s",
        args.ws_prefix,
    )

    # --------------------------------------------------------
    # Проверяем изображение
    # --------------------------------------------------------

    image_path = Path(
        args.image
    )

    if not image_path.is_file():

        LOG.error(
            "❌ Файл изображения не найден: %s",
            image_path,
        )

        return 1

    # --------------------------------------------------------
    # HTTP connector
    # --------------------------------------------------------

    connector = aiohttp.TCPConnector(
        ssl=False if args.insecure else None
    )

    timeout = aiohttp.ClientTimeout(
        total=args.http_timeout
    )

    uid = None
    access_token = None

    async with aiohttp.ClientSession(
        connector=connector,
        timeout=timeout,
    ) as session:

        try:

            # ====================================================
            # 1. Keycloak
            # ====================================================

            access_token = await get_access_token(
                session,
                args,
            )

            # ====================================================
            # 2. POST /api/v1/predict/new
            # ====================================================

            uid = await create_predict_session(
                session,
                args,
                access_token,
            )

            # ====================================================
            # 3. WebSocket
            # ====================================================

            success = await websocket_predict(
                args,
                access_token,
                uid,
            )

            if success:

                LOG.info("")
                LOG.info("=" * 70)
                LOG.info(
                    "✅ УСПЕШНО"
                )
                LOG.info("=" * 70)

            else:

                LOG.error("")
                LOG.error("=" * 70)
                LOG.error(
                    "❌ WEBSOCKET НЕДОСТУПЕН"
                )
                LOG.error("=" * 70)

                return 2

        except KeyboardInterrupt:

            LOG.warning(
                "Получен Ctrl+C"
            )

            return 130

        except Exception:

            LOG.exception(
                "❌ Критическая ошибка"
            )

            return 1

        finally:

            # ====================================================
            # 4. DELETE /api/v1/predict/{uid}
            # ====================================================

            if (
                uid
                and access_token
                and not args.no_delete
            ):

                await delete_predict_session(
                    session,
                    args,
                    access_token,
                    uid,
                )

            elif (
                uid
                and args.no_delete
            ):

                LOG.info(
                    "Session %s оставлена "
                    "(--no-delete)",
                    uid,
                )

    return 0


# ============================================================
# Entry point
# ============================================================

if __name__ == "__main__":

    try:
        exit_code = asyncio.run(
            main()
        )

    except KeyboardInterrupt:
        exit_code = 130

    sys.exit(exit_code)