from __future__ import annotations

import time

import aiohttp

from yuxi.channels.models import HealthStatus
from yuxi.utils.logging_config import logger


async def health_check_yuanbao(
    api_base: str,
    token: str,
    ws_connected: bool = False,
    session: aiohttp.ClientSession | None = None,
) -> HealthStatus:
    headers = {"Authorization": f"Bearer {token}"}

    async def _do_check(sess: aiohttp.ClientSession) -> HealthStatus:
        start = time.monotonic()
        async with sess.get(
            f"{api_base}/api/v1/bot/info",
            headers=headers,
            timeout=aiohttp.ClientTimeout(total=10),
        ) as resp:
            latency_ms = (time.monotonic() - start) * 1000

            if resp.status == 200:
                data = await resp.json()
                return HealthStatus(
                    status="healthy",
                    latency_ms=latency_ms,
                    metadata={
                        "bot_app_id": data.get("bot_app_id"),
                        "ws_connected": ws_connected,
                    },
                )
            elif resp.status == 401:
                return HealthStatus(
                    status="unhealthy",
                    last_error="Token expired or invalid",
                    metadata={"auth_status": "failed"},
                )
            else:
                return HealthStatus(
                    status="degraded",
                    latency_ms=latency_ms,
                    last_error=f"Bot info returned {resp.status}",
                )

    try:
        if session:
            return await _do_check(session)
        async with aiohttp.ClientSession() as new_session:
            return await _do_check(new_session)
    except Exception as e:
        logger.warning(f"[Yuanbao] Health check failed: {e}")
        return HealthStatus(
            status="unhealthy",
            last_error=str(e),
        )