"""TDengine 连接模块：taos-ws-py WebSocket 封装 + 指数退避重试。"""

import time
import logging
import taosws

logger = logging.getLogger(__name__)

_MAX_RETRY_DURATION = 1800      # 30 分钟
_MAX_BACKOFF = 300              # 5 分钟
_INSERT_BATCH_SIZE = 500


class ConnectionError(Exception):
    """TDengine 连接不可用时抛出。"""
    pass


class TDConnection:
    """TDengine WebSocket 连接封装。

    Args:
        url: WebSocket 连接字符串，默认 ws://root:taosdata@localhost:6041
    """

    def __init__(self, url: str = "ws://root:taosdata@localhost:6041"):
        self._url = url
        self._conn: taosws.Connection | None = None

    @property
    def conn(self) -> taosws.Connection:
        if self._conn is None:
            self._conn = self._connect_with_retry()
        return self._conn

    def _connect_with_retry(self) -> taosws.Connection:
        """指数退避重试连接 TDengine。

        初始退避 1s，倍增因子 2，最大间隔 300s，最长持续 1800s。
        每次重试前执行 SELECT SERVER_VERSION() 验证服务可用。

        Raises:
            ConnectionError: 30 分钟内仍不可用
        """
        backoff = 1.0
        started_at = time.monotonic()
        last_exc = None

        while True:
            elapsed = time.monotonic() - started_at
            if elapsed > _MAX_RETRY_DURATION:
                raise ConnectionError(
                    f"TDengine 连接失败，已重试 {elapsed:.0f}s（超过 30 分钟上限）"
                ) from last_exc

            try:
                conn = taosws.connect(self._url)
                conn.execute("SELECT SERVER_VERSION()")
                logger.info("TDengine 连接成功 (url=%s)", self._url)
                return conn
            except Exception as e:
                last_exc = e
                logger.warning(
                    "TDengine 连接失败 (attempt, backoff=%.1fs, elapsed=%.0fs): %s",
                    backoff, elapsed, e,
                )
                time.sleep(backoff)
                backoff = min(backoff * 2, _MAX_BACKOFF)

    def execute(self, sql: str) -> list[tuple]:
        """执行查询 SQL，返回结果行列表。"""
        logger.debug("execute: %s", sql[:120])
        result = self.conn.execute(sql)
        rows = []
        if result is not None:
            for row in result:
                rows.append(tuple(row))
        return rows

    def execute_no_result(self, sql: str) -> None:
        """执行非查询 SQL（INSERT / DDL）。"""
        logger.debug("execute_no_result: %s", sql[:120])
        self.conn.execute(sql)

    def close(self):
        """关闭连接。"""
        if self._conn is not None:
            self._conn.close()
            self._conn = None

    def is_healthy(self) -> bool:
        """检查连接健康状态。"""
        try:
            self.execute("SELECT SERVER_VERSION()")
            return True
        except Exception:
            return False

    def batch_insert(self, stable: str, rows: list[tuple[int, float, int]]) -> None:
        """批量 INSERT 数据到 subtable。

        每 500 条一个 batch，单次 execute_no_result()。

        Args:
            stable: supertable 名称（不含 ds_ 前缀）
            rows: [(ts_ms, value, label), ...]
        """
        total = len(rows)
        for i in range(0, total, _INSERT_BATCH_SIZE):
            batch = rows[i:i + _INSERT_BATCH_SIZE]
            values_clauses = []
            for ts, val, label in batch:
                values_clauses.append(f"({ts}, {val}, {label})")
            sql = f"INSERT INTO s_{stable} VALUES {' '.join(values_clauses)}"
            self.execute_no_result(sql)
            logger.debug("batch_insert: %d/%d rows", min(i + _INSERT_BATCH_SIZE, total), total)
        logger.info("batch_insert complete: %d rows into s_%s", total, stable)