"""
StarRocks Database Connection Utility
Based on chatbot-rsa pattern
"""

import asyncio
import logging
from typing import Any

import aiomysql
import pymysql
from pymysql.cursors import DictCursor

from config import (
    STARROCKS_DB,
    STARROCKS_HOST,
    STARROCKS_PASSWORD,
    STARROCKS_PORT,
    STARROCKS_USER,
)

logger = logging.getLogger(__name__)

__all__ = ["StarRocksConnection", "get_db_connection"]


class StarRocksConnectionManager:
    """
    Singleton Class quản lý StarRocks Connection.
    """

    def __init__(self):
        self._connection: StarRocksConnection | None = None

    def get_connection(self) -> "StarRocksConnection":
        """Lazy loading connection"""
        if self._connection is None:
            logger.info("🔧 [LAZY LOADING] Creating StarRocksConnection instance (first time)")
            self._connection = StarRocksConnection()
        return self._connection


# --- Singleton ---
_manager = StarRocksConnectionManager()
get_db_connection = _manager.get_connection


class StarRocksConnection:
    # Shared connection (Singleton-like behavior) for all instances
    _shared_conn = None

    def __init__(
        self,
        host: str | None = None,
        database: str | None = None,
        user: str | None = None,
        password: str | None = None,
        port: int | None = None,
    ):
        self.host = host or STARROCKS_HOST
        self.database = database or STARROCKS_DB
        self.user = user or STARROCKS_USER
        self.password = password or STARROCKS_PASSWORD
        self.port = port or STARROCKS_PORT
        # self.conn references the shared connection
        self.conn = None
        logger.info(f"✅ StarRocksConnection initialized: {self.host}:{self.port}")

    def connect(self):
        """
        Establish or reuse persistent connection.
        """
        # 1. Try to reuse existing shared connection
        if StarRocksConnection._shared_conn and StarRocksConnection._shared_conn.open:
            try:
                # Ping to check if alive, reconnect if needed
                StarRocksConnection._shared_conn.ping(reconnect=True)
                self.conn = StarRocksConnection._shared_conn
                return self.conn
            except Exception as e:
                logger.warning(f"⚠️ Connection lost, reconnecting: {e}")
                StarRocksConnection._shared_conn = None

        # 2. Create new connection if needed
        print(f"   [DB] 🔌 Đang kết nối StarRocks (New Session): {self.host}:{self.port}...")
        logger.info(f"🔌 Connecting to StarRocks at {self.host}:{self.port} (DB: {self.database})...")
        try:
            new_conn = pymysql.connect(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=self.database,
                charset="utf8mb4",
                cursorclass=DictCursor,
                connect_timeout=10,
                read_timeout=30,
                write_timeout=30,
            )
            print("   [DB] ✅ Kết nối thành công.")
            logger.info("✅ Connected to StarRocks")

            # Save to class variable
            StarRocksConnection._shared_conn = new_conn
            self.conn = new_conn

        except Exception as e:
            print(f"   [DB] ❌ Lỗi kết nối: {e!s}")
            logger.error(f"❌ Failed to connect to StarRocks: {e}")
            raise

        return self.conn

    def execute_query(self, query: str, params: tuple | None = None) -> list[dict[str, Any]]:
        # print("   [DB] 🚀 Bắt đầu truy vấn dữ liệu...")
        # (Reduced noise in logs)
        logger.info("🚀 Executing StarRocks Query (Persistent Conn).")

        conn = self.connect()
        try:
            with conn.cursor() as cursor:
                cursor.execute(query, params)
                results = cursor.fetchall()
                print(f"   [DB] ✅ Truy vấn xong. Lấy được {len(results)} dòng.")
                logger.info(f"📊 Query successful, returned {len(results)} rows")
                return [dict(row) for row in results]
        except Exception as e:
            print(f"   [DB] ❌ Lỗi truy vấn: {e!s}")
            logger.error(f"❌ StarRocks query error: {e}")
            # Incase of query error due to connection, invalidate it
            StarRocksConnection._shared_conn = None
            raise
        # FINALLY BLOCK REMOVED: Do NOT close connection

    # Async pool shared
    _shared_pool = None
    _pool_lock = asyncio.Lock()

    @classmethod
    async def clear_pool(cls):
        """Clear and close existing pool (force recreate fresh connections)"""
        async with cls._pool_lock:
            if cls._shared_pool is not None:
                logger.warning("🔄 Clearing StarRocks connection pool...")
                cls._shared_pool.close()
                await cls._shared_pool.wait_closed()
                cls._shared_pool = None
                logger.info("✅ Pool cleared successfully")

    async def get_pool(self):
        """
        Get or create shared async connection pool (Thread-safe singleton)
        """
        if StarRocksConnection._shared_pool is None:
            async with StarRocksConnection._pool_lock:
                # Double-check inside lock to prevent multiple pools
                if StarRocksConnection._shared_pool is None:
                    logger.info(f"🔌 Creating Async Pool to {self.host}:{self.port}...")
                    StarRocksConnection._shared_pool = await aiomysql.create_pool(
                            host=self.host,
                            port=self.port,
                            user=self.user,
                            password=self.password,
                            db=self.database,
                            charset="utf8mb4",
                            cursorclass=aiomysql.DictCursor,
                            minsize=5,        # ← Từ 10 → 5
                            maxsize=30,       # ← Từ 200 → 30 (QUAN TRỌNG!)
                            connect_timeout=10,
                            pool_recycle=1800,  # ← Từ 3600 → 1800 (30 phút)
                            autocommit=True,
                        )
        return StarRocksConnection._shared_pool

    async def execute_query_async(self, query: str, params: tuple | None = None) -> list[dict[str, Any]]:
        """
        Execute query asynchronously using aiomysql pool with Retry Logic.
        """
        max_retries = 3
        last_error = None

        for attempt in range(max_retries):
            try:
                pool = await self.get_pool()
                # logger.info(f"🚀 Executing Async Query (Attempt {attempt+1}).")

                # Tăng timeout lên 30s cho load test với 300 users
                conn = await asyncio.wait_for(pool.acquire(), timeout=30)
                try:
                    async with conn.cursor() as cursor:
                        await cursor.execute(query, params)
                        results = await cursor.fetchall()
                        # logger.info(f"📊 Async Query successful, returned {len(results)} rows")
                        return [dict(row) for row in results]
                finally:
                    pool.release(conn)

            except TimeoutError as e:
                last_error = e
                logger.warning(f"⏱️ Pool acquire timeout (Attempt {attempt + 1}/{max_retries})")
                # Timeout khi lấy connection → pool đầy, chờ rồi thử lại
                await asyncio.sleep(0.2 * (attempt + 1))
                continue
            except ConnectionAbortedError as e:
                last_error = e
                logger.warning(f"🔌 Connection aborted (Attempt {attempt + 1}/{max_retries}): {e}")
                # Connection bị abort → clear pool và thử lại với fresh connections
                if attempt < max_retries - 1:
                    await StarRocksConnection.clear_pool()
                    await asyncio.sleep(0.3)
                continue
            except Exception as e:
                last_error = e
                logger.warning(f"⚠️ StarRocks DB Error (Attempt {attempt + 1}/{max_retries}): {e}")
                if "Memory of process exceed limit" in str(e):
                    # Nếu StarRocks OOM, đợi một chút rồi thử lại
                    await asyncio.sleep(0.5 * (attempt + 1))
                    continue
                if "Disconnected" in str(e) or "Lost connection" in str(e) or "aborted" in str(e).lower():
                    # Nếu mất kết nối, clear pool và thử lại
                    if attempt < max_retries - 1:
                        await StarRocksConnection.clear_pool()
                        await asyncio.sleep(0.3)
                    continue
                # Các lỗi khác (cú pháp,...) thì raise luôn
                raise

        logger.error(f"❌ Failed after {max_retries} attempts: {last_error}")
        raise last_error

    def close(self):
        """Explicitly close if needed (e.g. app shutdown)"""
        if StarRocksConnection._shared_conn and StarRocksConnection._shared_conn.open:
            StarRocksConnection._shared_conn.close()
            StarRocksConnection._shared_conn = None
            self.conn = None
