"""
Wolfx 防灾免费 API
"""

import asyncio
import json
from typing import Literal, TypeVar, cast
import aiohttp
from aiosignal import Signal
from loguru import logger
from pydantic import BaseModel, RootModel
import pydantic

from konabot.common.appcontext import after_init


class ScEewReport(BaseModel):
    """
    四川地震局报文
    """

    ID: str
    "EEW 发报 ID"

    EventID: str
    "EEW 发报事件 ID"

    ReportTime: str
    "EEW 发报时间(UTC+8)"

    ReportNum: int
    "EEW 发报数"

    OriginTime: str
    "发震时间(UTC+8)"

    HypoCenter: str
    "震源地"

    Latitude: float
    "震源地纬度"

    Longitude: float
    "震源地经度"

    Magnitude: float
    "震级"

    Depth: float | None
    "震源深度"

    MaxIntensity: float
    "最大烈度"


class CencEewReport(BaseModel):
    """
    中国地震台网报文
    """

    ID: str
    "EEW 发报 ID"

    EventID: str
    "EEW 发报事件 ID"

    ReportTime: str
    "EEW 发报时间(UTC+8)"

    ReportNum: int
    "EEW 发报数"

    OriginTime: str
    "发震时间(UTC+8)"

    HypoCenter: str
    "震源地"

    Latitude: float
    "震源地纬度"

    Longitude: float
    "震源地经度"

    Magnitude: float
    "震级"

    Depth: float | None
    "震源深度"

    MaxIntensity: float
    "最大烈度"


class CencEqReport(BaseModel):
    type: str
    "报告类型"

    EventID: str
    "事件 ID"

    time: str
    "UTC+8 格式的地震发生时间"

    location: str
    "地震发生位置"

    magnitude: str
    "震级"

    depth: str
    "地震深度"

    latitude: str
    "纬度"

    longtitude: str
    "经度"

    intensity: str
    "烈度"


class CencEqlist(RootModel):
    root: dict[str, CencEqReport]


class WolfxWebSocket:
    def __init__(self, url: str) -> None:
        self.url = url
        self.signal: Signal[bytes] = Signal(self)
        self._running = False
        self._task: asyncio.Task | None = None
        self._session: aiohttp.ClientSession | None = None
        self._ws: aiohttp.ClientWebSocketResponse | None = None

    @property
    def session(self) -> aiohttp.ClientSession:  # pragma: no cover
        assert self._session is not None
        return self._session

    async def start(self):  # pragma: no cover
        if self._running:
            return
        self._running = True
        self._session = aiohttp.ClientSession()
        self._task = asyncio.create_task(self._run())
        self.signal.freeze()

    async def stop(self):  # pragma: no cover
        self._running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            if self._session:
                await self._session.close()

    async def _run(self):  # pragma: no cover
        retry_delay = 1

        while self._running:
            try:
                async with self.session.ws_connect(self.url) as ws:
                    self._ws = ws
                    logger.info(f"Wolfx API 服务连接上了 {self.url} 的 WebSocket")
                    async for msg in ws:
                        if msg.type == aiohttp.WSMsgType.TEXT:
                            await self.handle(cast(str, msg.data).encode())
                        elif msg.type == aiohttp.WSMsgType.BINARY:
                            await self.handle(cast(bytes, msg.data))
                        elif msg.type == aiohttp.WSMsgType.CLOSED:
                            break
                        elif msg.type == aiohttp.WSMsgType.ERROR:
                            break
            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
                logger.warning("连接 WebSocket 时发生错误")
                logger.exception(e)
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error("Wolfx API 发生未知错误")
                logger.exception(e)
            self._ws = None

            if self._running:
                logger.info(f"Wolfx API 准备断线重连 {self.url}")
                await asyncio.sleep(retry_delay)
                retry_delay = min(retry_delay * 2, 60)

    async def handle(self, data: bytes):
        try:
            obj = json.loads(data)
        except json.JSONDecodeError as e:
            logger.warning("解析 Wolfs API 时出错")
            logger.exception(e)
            return

        if obj.get("type") == "heartbeat" or obj.get("type") == "pong":
            logger.debug(f"Wolfx API 收到了来自 {self.url} 的心跳: {obj}")
        else:
            await self.signal.send(data)


T = TypeVar("T", bound=BaseModel)


class WolfxAPIService:
    sc_eew: Signal[ScEewReport]
    "四川地震局地震速报"

    cenc_eew: Signal[CencEewReport]
    "中国地震台网地震速报"

    cenc_eqlist: Signal[CencEqReport]
    "中国地震台网地震信息发布"

    def __init__(self) -> None:
        self.sc_eew = Signal(self)
        self._sc_eew_ws = WolfxWebSocket("wss://ws-api.wolfx.jp/sc_eew")
        WolfxAPIService.bind(self.sc_eew, self._sc_eew_ws, ScEewReport)

        self.cenc_eew = Signal(self)
        self._cenc_eew_ws = WolfxWebSocket("wss://ws-api.wolfx.jp/cenc_eew")
        WolfxAPIService.bind(self.cenc_eew, self._cenc_eew_ws, CencEewReport)

        self.cenc_eqlist = Signal(self)
        self._cenc_eqlist_ws = WolfxWebSocket("wss://ws-api.wolfx.jp/cenc_eqlist")
        WolfxAPIService.bind(self.cenc_eqlist, self._cenc_eqlist_ws, CencEqReport)

    @staticmethod
    def bind(signal: Signal[T], ws: WolfxWebSocket, t: type[T]):
        @ws.signal.append
        async def _(data: bytes):
            try:
                obj = t.model_validate_json(data)
                logger.info(f"接收到来自 Wolfx API 的信息：{data}")
                await signal.send(obj)
            except pydantic.ValidationError as e:
                logger.warning(f"解析 Wolfx API 时出错 URL={ws.url}")
                logger.error(e)

    async def start(self):  # pragma: no cover
        self.cenc_eew.freeze()
        self.sc_eew.freeze()
        self.cenc_eqlist.freeze()
        async with asyncio.TaskGroup() as task_group:
            if len(self.cenc_eew) > 0:
                task_group.create_task(self._cenc_eew_ws.start())

            if len(self.sc_eew) > 0:
                task_group.create_task(self._sc_eew_ws.start())

            if len(self.cenc_eqlist) > 0:
                task_group.create_task(self._cenc_eqlist_ws.start())

    async def stop(self):  # pragma: no cover
        async with asyncio.TaskGroup() as task_group:
            task_group.create_task(self._cenc_eew_ws.stop())
            task_group.create_task(self._sc_eew_ws.stop())
            task_group.create_task(self._cenc_eqlist_ws.stop())


wolfx_api = WolfxAPIService()


@after_init
def init():  # pragma: no cover
    import nonebot

    driver = nonebot.get_driver()

    @driver.on_startup
    async def _():
        await wolfx_api.start()

    @driver.on_shutdown
    async def _():
        await wolfx_api.stop()
