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

import asyncio
import logging
import os
import threading
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:
    # Use thread-local storage to ensure thread-safety across FastAPI sync workers
    _thread_local = threading.local()

    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 thread-local connection
        conn = getattr(StarRocksConnection._thread_local, 'conn', None)
        if conn and conn.open:
            try:
                # Ping to check if alive, reconnect if needed
                conn.ping(reconnect=True)
                self.conn = conn
                return self.conn
            except Exception as e:
                logger.warning(f"⚠️ Connection lost, reconnecting: {e}")
                StarRocksConnection._thread_local.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 thread-local variable
            StarRocksConnection._thread_local.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]]:
        from config import USE_LOCAL_SQLITE
        if USE_LOCAL_SQLITE:
            from .sqlite_mock import MockCursor
            mc = MockCursor()
            mc.execute(query, params)
            # Convert MockRow (tuple) → dict so callers can mutate via p["key"] = val
            return [dict(zip(row.keys(), row)) for row in mc.fetchall()]

        # 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._thread_local.conn = None
            raise
        # FINALLY BLOCK REMOVED: Do NOT close connection

    # Async pool shared
    _shared_pool = None
    _pool_lock: asyncio.Lock | None = None

    @classmethod
    async def clear_pool(cls):
        """Clear and close existing pool (force recreate fresh connections)"""
        if cls._pool_lock is None:
            cls._pool_lock = asyncio.Lock()
        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)
        Optimized for cosine similarity queries (~200ms)
        """
        if StarRocksConnection._pool_lock is None:
            StarRocksConnection._pool_lock = asyncio.Lock()
            
        if StarRocksConnection._shared_pool is None:
            async with StarRocksConnection._pool_lock:
                if StarRocksConnection._shared_pool is None:
                    logger.info(f"🔌 Creating Async Pool to {self.host}:{self.port}...")
                    minsize = int(os.getenv("STARROCKS_POOL_MINSIZE", "2"))
                    maxsize = int(os.getenv("STARROCKS_POOL_MAXSIZE", "20"))
                    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=minsize,  # Giảm minsize để đỡ tốn tài nguyên idle
                        maxsize=maxsize,
                        connect_timeout=10,
                        # --- CHỈNH SỬA QUAN TRỌNG Ở ĐÂY ---
                        pool_recycle=3600,  # Recycle sau 4 phút rưỡi (tránh timeout 5 phút của Windows/Firewall)
                        # ----------------------------------
                        autocommit=True,
                    )
                    logger.info("✅ Pool created successfully with recycle=280s")
        return StarRocksConnection._shared_pool

    async def execute_query_async(self, query: str, params: tuple | None = None) -> list[dict[str, Any]]:
        """
        Execute query asynchronously with AUTO-RECONNECT (Fix lỗi 10053/2006).
        """
        from config import USE_LOCAL_SQLITE
        if USE_LOCAL_SQLITE:
            from .sqlite_mock import MockCursor
            mc = MockCursor()
            mc.execute(query, params)
            # Convert MockRow (tuple) → dict so callers can mutate via p["key"] = val
            return [dict(zip(row.keys(), row)) for row in mc.fetchall()]

        max_retries = 3

        for attempt in range(max_retries):
            pool = None
            conn = None
            try:
                pool = await self.get_pool()
                conn = await asyncio.wait_for(pool.acquire(), timeout=19)

                async with conn.cursor() as cursor:
                    await cursor.execute(query, params)
                    results = await cursor.fetchall()
                    return [dict(row) for row in results]

            # --- SỬA ĐOẠN CATCH ERROR RỘNG HƠN ---
            except (
                TimeoutError,
                pymysql.err.OperationalError,
                pymysql.err.InterfaceError,
                ConnectionError,
                OSError,
            ) as e:
                error_msg = str(e).lower()
                error_code = e.args[0] if e.args else 0

                logger.warning(f"⚠️ DB Error (Attempt {attempt + 1}/{max_retries}): {e}")

                # Danh sách mã lỗi MySQL phổ biến khi mất kết nối
                mysql_conn_codes = [2006, 2013, 2014, 2003, 10053, 10054, 10060, 10061]

                # Điều kiện Retry:
                # 1. Mã lỗi nằm trong list
                # 2. Hoặc là lỗi hệ thống mạng (ConnectionError)
                # 3. Hoặc thông báo lỗi chứa từ khóa nhạy cảm
                is_conn_error = (
                    error_code in mysql_conn_codes
                    or isinstance(e, (ConnectionError, BrokenPipeError, ConnectionResetError))
                    or "abort" in error_msg
                    or "closed" in error_msg
                    or "reset" in error_msg
                    or "pipe" in error_msg
                )

                if is_conn_error:
                    logger.info("♻️ Connection dead. Clearing pool and retrying...")
                    await StarRocksConnection.clear_pool()
                    await asyncio.sleep(0.5)
                    continue  # RETRY NGAY

                # Nếu là lỗi cú pháp SQL (ProgrammingError) thì raise luôn, không retry
                raise e
            # --------------------------------------

            except Exception as e:
                logger.error(f"❌ Unexpected DB Error: {e}")
                raise e

            finally:
                if pool and conn:
                    try:
                        pool.release(conn)
                    except Exception:
                        pass

        raise Exception("Failed to execute query after retries.")

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