import asyncio
import hashlib
import json
import os
import time
from typing import Any, List, Optional, Tuple, Union
from urllib.parse import urlparse

from dotenv import find_dotenv
from psycopg import AsyncConnection
from psycopg.errors import UniqueViolation
from psycopg.pq import TransactionStatus
from pydantic import PostgresDsn
from pydantic_settings import BaseSettings, SettingsConfigDict

from common.database.exceptions import UniqueConstraintError
from common.logging import get_logger, span

logger = get_logger(__name__)


class _PooledConnection:
    """Async context manager proxy that returns connections to the pool."""

    def __init__(self, owner: "Db", conn: AsyncConnection):
        self._owner = owner
        self._conn = conn
        self._released = False

    def __getattr__(self, item):
        return getattr(self._conn, item)

    async def __aenter__(self) -> "_PooledConnection":
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self._release()

    async def close(self) -> None:
        await self._release()

    def _get_lock(self) -> asyncio.Lock:
        """Get the lock for this connection to prevent concurrent cursor operations.

        The lock is created eagerly when the connection is established, so we can
        retrieve it without any synchronization overhead.

        Uses asyncio.Lock for native async/await semantics. All database operations
        are async, so cross-thread synchronization is not needed.

        Returns:
            asyncio.Lock: Connection-specific lock
        """
        conn_id = id(self._conn)
        # Lock is created eagerly in _create_connection, so we can retrieve it directly
        # without needing synchronization
        return self._owner._connection_locks[conn_id]

    async def _release(self) -> None:
        if self._released:
            return
        self._released = True
        conn = self._conn
        self._conn = None
        if conn is None:
            return
        try:
            if not conn.closed:
                try:
                    await conn.rollback()
                except Exception:
                    # Ignore rollback errors; connection will be recycled/closed below.
                    pass
        finally:
            await self._owner._release_connection(conn)


class DBConfig(BaseSettings):
    """
    Database configuration settings loaded from .env file

    Attributes:
    -----------
    PG_DSN: str
        Postgres connection string
    """

    pg_dsn: Optional[PostgresDsn] = None
    db_secret: Optional[str] = None
    db_pool_max_size: int = 50  # Increased from 30 to handle high-load scenarios (100+ concurrent streams + cleanup tasks)
    # ~8 instances × 50 = 400 connections (with RDS Proxy multiplexing to ~80-100 actual DB connections)

    # RDS Proxy support for connection multiplexing
    db_use_proxy: bool = False
    db_proxy_endpoint: Optional[str] = None
    db_proxy_port: int = 5432

    # Slow query detection threshold in milliseconds
    db_slow_query_threshold_ms: int = 1000  # Log queries taking longer than 1 second

    model_config = SettingsConfigDict(
        env_file=(find_dotenv(), "dev.env"), env_file_encoding="utf-8", extra="ignore"
    )


class Db:
    """
    Asynchronous Database connection class for executing queries with direct connections.
    """

    def __init__(self, config: DBConfig) -> None:
        self.config = config
        self.secret: dict[str, Any] = {}

        if self.config.pg_dsn is None and self.config.db_secret is None:
            raise ValueError("Either pg_dsn or db_secret must be provided.")

        if self.config.db_secret is not None:
            db_secret = json.loads(self.config.db_secret)

            # Override with RDS Proxy endpoint if enabled
            if self.config.db_use_proxy and self.config.db_proxy_endpoint:
                host = self.config.db_proxy_endpoint
                port = self.config.db_proxy_port
                logger.info(
                    "Using RDS Proxy for database connections",
                    extra={
                        "proxy_endpoint": host,
                        "proxy_port": port,
                    },
                )
            else:
                host = db_secret["host"]
                port = int(db_secret["port"])  # Convert port string to int

            username = db_secret["username"]
            password = db_secret["password"]
            dbname = db_secret["dbname"]
            self.secret = {
                "host": host,
                "port": str(port),
                "username": username,
                "password": password,
                "dbname": dbname,
            }
            self.config.pg_dsn = PostgresDsn.build(
                scheme="postgresql+psycopg",
                username=username,
                password=password,
                host=host,
                port=port,
                path=f"{dbname}",
            )
        elif self.config.pg_dsn is not None:
            # Extract fields for diagnostics/logging when only DSN is supplied
            dsn_obj = self.config.pg_dsn
            parsed = urlparse(str(dsn_obj))
            host = parsed.hostname or "localhost"
            port = parsed.port or 5432
            username = parsed.username or ""
            password = parsed.password or ""
            dbname = (parsed.path or "").lstrip("/") or "postgres"
            self.secret = {
                "host": host,
                "port": str(port),
                "username": username,
                "password": password,
                "dbname": dbname,
            }

        self.dsn = str(self.config.pg_dsn)
        self._conninfo = (
            f"host={self.secret['host']} "
            f"port={self.secret['port']} "
            f"user={self.secret['username']} "
            f"password={self.secret['password']} "
            f"dbname={self.secret['dbname']} "
            f"connect_timeout=30"  # Increased from 10 to reduce timeout errors during high load
        )
        self._pool_lock = asyncio.Lock()
        self._pool_queue: asyncio.Queue[AsyncConnection] = asyncio.Queue(
            maxsize=max(self.config.db_pool_max_size, 1)
        )
        self._active_connections = 0
        self._connection_max_age = 600  # 10 minutes - recycle connections
        self._connection_timestamps: dict[
            int, float
        ] = {}  # Track connection creation time
        self._connection_locks: dict[
            int, asyncio.Lock
        ] = {}  # Per-connection asyncio locks (created eagerly) to prevent "another command in progress" errors
        self._slow_query_log_cooldown_seconds = max(
            0, int(os.getenv("TLDR_DB_SLOW_QUERY_COOLDOWN_SECONDS", "0"))
        )
        self._slow_query_last_logged: dict[str, float] = {}
        self._slow_query_log_lock = asyncio.Lock()

    async def get_connection(self) -> _PooledConnection:
        """
        Get a new database connection.

        Returns:
            _PooledConnection: A pooled database connection wrapper

        Raises:
            Exception: If connection fails
        """
        connection_attributes = {
            "host": self.secret.get("host", "unknown"),
            "port": self.secret.get("port", "unknown"),
            "dbname": self.secret.get("dbname", "unknown"),
            "user": self.secret.get("username", "unknown"),
        }

        with span(
            logger,
            "get_database_connection",
            connection_attributes,
            log_entry=False,
            log_exit=False,
            low_level=True,
        ):
            try:
                conn = await self._acquire_connection()
                return _PooledConnection(self, conn)
            except asyncio.CancelledError:
                logger.warning(
                    "Database connection attempt was cancelled",
                    extra={"host": connection_attributes["host"]},
                )
                raise
            except Exception as e:
                logger.error(
                    "Failed to retrieve pooled connection",
                    extra={"host": connection_attributes["host"], "error": str(e)},
                )
                raise

    async def _acquire_connection(self) -> AsyncConnection:
        import time

        current_time = time.time()

        # Fast path: reuse idle connection from queue
        while True:
            try:
                conn = self._pool_queue.get_nowait()
            except asyncio.QueueEmpty:
                break

            # Check if connection is closed or too old
            conn_id = id(conn)
            if conn.closed:
                async with self._pool_lock:
                    self._active_connections = max(0, self._active_connections - 1)
                    self._connection_timestamps.pop(conn_id, None)
                    self._connection_locks.pop(conn_id, None)  # Clean up lock
                continue

            # Recycle connections older than max_age (10 minutes)
            conn_age = current_time - self._connection_timestamps.get(
                conn_id, current_time
            )
            if conn_age > self._connection_max_age:
                try:
                    await conn.close()
                except Exception:
                    pass
                async with self._pool_lock:
                    self._active_connections = max(0, self._active_connections - 1)
                    self._connection_timestamps.pop(conn_id, None)
                    self._connection_locks.pop(conn_id, None)  # Clean up lock
                continue

            # Validate connection health before returning
            if not await self._validate_connection(conn):
                logger.warning(
                    "Connection failed health check, closing and trying next",
                    extra={"conn_id": conn_id},
                )
                try:
                    await conn.close()
                except Exception:
                    pass
                async with self._pool_lock:
                    self._active_connections = max(0, self._active_connections - 1)
                    self._connection_timestamps.pop(conn_id, None)
                    self._connection_locks.pop(conn_id, None)
                continue

            return conn

        # No idle connection available; attempt to create a new one if under limit
        async with self._pool_lock:
            if self._active_connections < self.config.db_pool_max_size:
                conn = await self._create_connection()
                self._active_connections += 1
                # Track connection creation time
                self._connection_timestamps[id(conn)] = current_time
                return conn

        # Pool is at max capacity; wait for a connection to be released
        # Add timeout to prevent indefinite waits when pool is exhausted
        try:
            async with asyncio.timeout(30):  # 30 second timeout
                while True:
                    conn = await self._pool_queue.get()
                    if conn.closed:
                        async with self._pool_lock:
                            self._active_connections = max(
                                0, self._active_connections - 1
                            )
                            self._connection_timestamps.pop(id(conn), None)
                            self._connection_locks.pop(id(conn), None)  # Clean up lock
                        continue

                    # Check connection age when acquired from queue
                    conn_id = id(conn)
                    conn_age = current_time - self._connection_timestamps.get(
                        conn_id, current_time
                    )
                    if conn_age > self._connection_max_age:
                        try:
                            await conn.close()
                        except Exception:
                            pass
                        async with self._pool_lock:
                            self._active_connections = max(
                                0, self._active_connections - 1
                            )
                            self._connection_timestamps.pop(conn_id, None)
                            self._connection_locks.pop(conn_id, None)  # Clean up lock
                        continue

                    # Validate connection health before returning
                    if not await self._validate_connection(conn):
                        logger.warning(
                            "Connection failed health check during wait, closing and trying next",
                            extra={"conn_id": conn_id},
                        )
                        try:
                            await conn.close()
                        except Exception:
                            pass
                        async with self._pool_lock:
                            self._active_connections = max(
                                0, self._active_connections - 1
                            )
                            self._connection_timestamps.pop(conn_id, None)
                            self._connection_locks.pop(conn_id, None)
                        continue

                    return conn
        except asyncio.TimeoutError:
            raise RuntimeError(
                f"Database connection pool exhausted: timeout after 30s waiting for connection "
                f"(pool_size={self.config.db_pool_max_size}, active={self._active_connections})"
            )

    async def _create_connection(self) -> AsyncConnection:
        conn = await AsyncConnection.connect(self._conninfo)
        # Eagerly create asyncio.Lock for this connection to prevent concurrent operations
        # Using asyncio.Lock provides better performance than threading.Lock + asyncio.to_thread
        self._connection_locks[id(conn)] = asyncio.Lock()
        return conn

    async def _validate_connection(self, conn: AsyncConnection) -> bool:
        """Validate connection health with a quick health check query.

        Args:
            conn: Connection to validate

        Returns:
            bool: True if connection is healthy, False otherwise
        """
        try:
            # Quick health check query
            async with conn.cursor() as cur:
                await cur.execute(b"SELECT 1")
                result = await cur.fetchone()
                return result is not None and result[0] == 1
        except Exception as e:
            logger.warning(
                "Connection health check failed",
                extra={"conn_id": id(conn), "error": str(e)},
            )
            return False

    async def _release_connection(self, conn: AsyncConnection) -> None:
        conn_id = id(conn)
        if conn.closed:
            async with self._pool_lock:
                self._active_connections = max(0, self._active_connections - 1)
                self._connection_timestamps.pop(conn_id, None)
                self._connection_locks.pop(conn_id, None)  # Clean up lock
            return

        # CRITICAL: Reset connection state before returning to pool
        # This prevents "another command is already in progress" errors
        try:
            # Check and reset transaction state
            trans_status = conn.info.transaction_status

            if trans_status == TransactionStatus.INTRANS:
                # Connection has an open transaction - rollback to clean state
                await conn.rollback()
                logger.warning(
                    "Connection had open transaction, rolled back before returning to pool",
                    extra={"conn_id": conn_id, "transaction_status": trans_status.name},
                )

            elif trans_status == TransactionStatus.ACTIVE:
                # Connection has an active query - this shouldn't happen with locks
                # but handle defensively
                logger.error(
                    "Connection has active query during release - closing connection",
                    extra={"conn_id": conn_id, "transaction_status": trans_status.name},
                )
                # Close this connection instead of returning to pool
                await conn.close()
                async with self._pool_lock:
                    self._active_connections = max(0, self._active_connections - 1)
                    self._connection_timestamps.pop(conn_id, None)
                    self._connection_locks.pop(conn_id, None)
                return

            elif trans_status == TransactionStatus.INERROR:
                # Transaction is in error state - rollback to clean it
                await conn.rollback()
                logger.warning(
                    "Connection in error state, rolled back before returning to pool",
                    extra={"conn_id": conn_id, "transaction_status": trans_status.name},
                )

        except Exception as e:
            # If state reset fails, close the connection instead of returning to pool
            logger.error(
                "Failed to reset connection state, closing instead of returning to pool",
                extra={"conn_id": conn_id, "error": str(e)},
            )
            try:
                await conn.close()
            except Exception:
                pass
            async with self._pool_lock:
                self._active_connections = max(0, self._active_connections - 1)
                self._connection_timestamps.pop(conn_id, None)
                self._connection_locks.pop(conn_id, None)
            return

        # Put cleaned connection back in pool
        try:
            self._pool_queue.put_nowait(conn)
        except asyncio.QueueFull:
            await conn.close()
            async with self._pool_lock:
                self._active_connections = max(0, self._active_connections - 1)
                self._connection_timestamps.pop(conn_id, None)
                self._connection_locks.pop(conn_id, None)  # Clean up lock

    def _extract_query_type(self, query: str) -> str:
        """
        Extract the query type (SELECT, INSERT, etc.) from a SQL query.

        Args:
            query: The SQL query

        Returns:
            str: The query type in uppercase, or "UNKNOWN" if it can't be determined
        """
        return query.strip().split(" ")[0].upper() if query.strip() else "UNKNOWN"

    def _encode_query(self, query: str) -> bytes:
        """
        Encode a SQL query to UTF-8 bytes.

        Args:
            query: The SQL query

        Returns:
            bytes: The query encoded as UTF-8 bytes
        """
        return query.encode("utf-8")

    async def _should_log_slow_query(self, query_preview: str) -> bool:
        """Apply per-query cooldown so repeated slow statements do not flood logs."""

        if self._slow_query_log_cooldown_seconds <= 0:
            return True

        query_hash = hashlib.sha256(query_preview.encode("utf-8")).hexdigest()
        now = time.monotonic()

        async with self._slow_query_log_lock:
            last_logged = self._slow_query_last_logged.get(query_hash)
            if (
                last_logged is not None
                and now - last_logged < self._slow_query_log_cooldown_seconds
            ):
                return False

            self._slow_query_last_logged[query_hash] = now
            if len(self._slow_query_last_logged) > 5000:
                expiry = now - (self._slow_query_log_cooldown_seconds * 2)
                self._slow_query_last_logged = {
                    key: timestamp
                    for key, timestamp in self._slow_query_last_logged.items()
                    if timestamp >= expiry
                }
            return True

    async def _log_slow_query_warning(
        self,
        *,
        query_type: str,
        duration_ms: float,
        query_preview: str,
        extra_fields: Optional[dict[str, Any]] = None,
    ) -> None:
        """Emit a slow-query warning with cooldown deduplication."""

        if not await self._should_log_slow_query(query_preview):
            return

        log_extra = {
            "query_type": query_type,
            "duration_ms": round(duration_ms, 2),
            "threshold_ms": self.config.db_slow_query_threshold_ms,
            "query_preview": query_preview,
        }
        if extra_fields:
            log_extra.update(extra_fields)

        logger.warning("Slow query detected", extra=log_extra)

    async def execute(
        self, query: str, params: Optional[Union[Tuple, List]] = None
    ) -> Optional[List[Tuple[Any, ...]]]:
        """
        Execute a SQL query and optionally fetch all results.

        Args:
            query: The SQL query to execute
            params: Optional parameters for the query

        Returns:
            Optional[List[Tuple[Any, ...]]]: Query results if the query returns data, None otherwise

        Raises:
            Exception: If query execution fails
        """
        import time

        # Extract query type for logging (SELECT, INSERT, etc.)
        query_type = self._extract_query_type(query)

        # Create attributes for span without including potentially sensitive params
        span_attributes = {"query_type": query_type, "has_params": params is not None}

        start_time = time.perf_counter()

        with span(
            logger,
            "execute_query",
            span_attributes,
            log_entry=False,
            log_exit=False,
            low_level=True,
        ):
            async with await self.get_connection() as conn:
                # Acquire connection-level lock to prevent concurrent cursor operations
                lock = conn._get_lock()
                async with lock:
                    async with conn.cursor() as cur:
                        try:
                            await cur.execute(self._encode_query(query), params)
                        except Exception as e:
                            logger.error(
                                f"Error executing {query_type} query: {str(e)}"
                            )
                            raise

                        if cur.description:  # If the query returns data
                            results = await cur.fetchall()

                            # Check for slow query
                            duration_ms = (time.perf_counter() - start_time) * 1000
                            if duration_ms > self.config.db_slow_query_threshold_ms:
                                # Sanitize query for logging (remove newlines, limit length)
                                query_preview = query.replace("\n", " ").strip()[:200]
                                await self._log_slow_query_warning(
                                    query_type=query_type,
                                    duration_ms=duration_ms,
                                    query_preview=query_preview,
                                    extra_fields={
                                        "row_count": len(results) if results else 0
                                    },
                                )

                            return results

                        await conn.commit()

            return None

    async def fetch_one(
        self, query: str, params: Optional[Union[Tuple, List]] = None
    ) -> Optional[Tuple[Any, ...]]:
        """
        Execute a SQL query and fetch a single result.

        Args:
            query: The SQL query to execute
            params: Optional parameters for the query

        Returns:
            Optional[Tuple[Any, ...]]: A single row result or None if no results

        Raises:
            Exception: If query execution fails
        """
        import time

        # Extract query type for logging (SELECT, INSERT, etc.)
        query_type = self._extract_query_type(query)

        # Create attributes for span without including potentially sensitive params
        span_attributes = {"query_type": query_type, "has_params": params is not None}

        start_time = time.perf_counter()

        with span(
            logger,
            "fetch_one",
            span_attributes,
            log_entry=False,
            log_exit=False,
            low_level=True,
        ):
            async with await self.get_connection() as conn:
                # Acquire connection-level lock to prevent concurrent cursor operations
                lock = conn._get_lock()
                async with lock:
                    async with conn.cursor() as cur:
                        try:
                            await cur.execute(self._encode_query(query), params)
                        except Exception as e:
                            logger.error(
                                f"Error executing query to fetch one row: {str(e)}"
                            )
                            raise

                        result = await cur.fetchone()

                        # Check for slow query
                        duration_ms = (time.perf_counter() - start_time) * 1000
                        if duration_ms > self.config.db_slow_query_threshold_ms:
                            # Sanitize query for logging (remove newlines, limit length)
                            query_preview = query.replace("\n", " ").strip()[:200]
                            await self._log_slow_query_warning(
                                query_type=query_type,
                                duration_ms=duration_ms,
                                query_preview=query_preview,
                                extra_fields={"has_result": result is not None},
                            )

                        return result

    async def fetch_all(
        self, query: str, params: Optional[Union[Tuple, List]] = None
    ) -> Optional[List[Tuple[Any, ...]]]:
        """
        Execute a SQL query and fetch all results.

        Args:
            query: The SQL query to execute
            params: Optional parameters for the query

        Returns:
            Optional[List[Tuple[Any, ...]]]: All rows returned by the query or empty list if no results

        Raises:
            Exception: If query execution fails
        """
        import time

        # Extract query type for logging (SELECT, INSERT, etc.)
        query_type = self._extract_query_type(query)

        # Create attributes for span without including potentially sensitive params
        span_attributes = {"query_type": query_type, "has_params": params is not None}

        start_time = time.perf_counter()

        with span(
            logger,
            "fetch_all",
            span_attributes,
            log_entry=False,
            log_exit=False,
            low_level=True,
        ):
            async with await self.get_connection() as conn:
                # Acquire connection-level lock to prevent concurrent cursor operations
                lock = conn._get_lock()
                async with lock:
                    async with conn.cursor() as cur:
                        try:
                            await cur.execute(self._encode_query(query), params)
                        except Exception as e:
                            logger.error(
                                f"Error executing query to fetch all rows: {str(e)}"
                            )
                            raise

                        results = await cur.fetchall()

                        # Check for slow query
                        duration_ms = (time.perf_counter() - start_time) * 1000
                        if duration_ms > self.config.db_slow_query_threshold_ms:
                            # Sanitize query for logging (remove newlines, limit length)
                            query_preview = query.replace("\n", " ").strip()[:200]
                            await self._log_slow_query_warning(
                                query_type=query_type,
                                duration_ms=duration_ms,
                                query_preview=query_preview,
                                extra_fields={
                                    "row_count": len(results) if results else 0
                                },
                            )

                        return results

    async def execute_commit(
        self, query: str, params: Optional[Union[Tuple, List]] = None
    ) -> Optional[int]:
        """
        Execute a SQL command that modifies the database and commit the transaction.

        Args:
            query: The SQL query to execute
            params: Optional parameters for the query

        Returns:
            Optional[int]: Number of affected rows when available (e.g. UPDATE/DELETE), else None.

        Raises:
            Exception: If query execution or commit fails
        """
        import time

        # Extract query type for logging (INSERT, UPDATE, DELETE, etc.)
        query_type = self._extract_query_type(query)

        # Create attributes for span without including potentially sensitive params
        span_attributes = {"query_type": query_type, "has_params": params is not None}

        start_time = time.perf_counter()

        with span(
            logger,
            "execute_commit",
            span_attributes,
            log_entry=False,
            log_exit=False,
            low_level=True,
        ):
            async with await self.get_connection() as conn:
                # Acquire connection-level lock to prevent concurrent cursor operations
                lock = conn._get_lock()
                async with lock:
                    affected_rows: Optional[int] = None
                    async with conn.cursor() as cur:
                        try:
                            await cur.execute(self._encode_query(query), params)
                            # Capture affected row count when available
                            affected_rows = getattr(cur, "rowcount", None)
                        except Exception as e:
                            logger.error(
                                f"Error executing {query_type} query: {str(e)}"
                            )
                            raise

                    try:
                        await conn.commit()
                    except Exception as e:
                        logger.error(f"Error committing transaction: {str(e)}")
                        raise

                    # Check for slow query
                    duration_ms = (time.perf_counter() - start_time) * 1000
                    if duration_ms > self.config.db_slow_query_threshold_ms:
                        # Sanitize query for logging (remove newlines, limit length)
                        query_preview = query.replace("\n", " ").strip()[:200]
                        await self._log_slow_query_warning(
                            query_type=query_type,
                            duration_ms=duration_ms,
                            query_preview=query_preview,
                            extra_fields={"affected_rows": affected_rows},
                        )

                    return affected_rows

    async def insert(
        self,
        table: str,
        columns: List[str],
        values: Union[Any, Tuple[Any, ...]],
        on_conflict: str = "ON CONFLICT DO NOTHING",
    ) -> Optional[int]:
        """
        Insert data into a specified table and return the id.

        Args:
            table: The name of the table to insert into
            columns: List of column names
            values: Values to insert
            on_conflict: SQL clause for handling conflicts (default: "ON CONFLICT DO NOTHING")

        Returns:
            Optional[int]: The ID of the inserted record or None if insert failed

        Raises:
            UniqueConstraintError: If a unique constraint is violated
            Exception: If insert fails for any other reason
        """
        # Create attributes for span without including potentially sensitive values
        span_attributes = {
            "table": table,
            "column_count": len(columns),
            "operation": "INSERT",
        }

        with span(
            logger,
            "insert_data",
            span_attributes,
            log_entry=False,
            log_exit=False,
        ):
            placeholders = ", ".join(["%s"] * len(columns))
            columns_formatted = ", ".join(columns)
            query = f"INSERT INTO {table} ({columns_formatted}) VALUES ({placeholders}) RETURNING id;"

            async with await self.get_connection() as conn:
                # Acquire connection-level lock to prevent concurrent cursor operations
                lock = conn._get_lock()
                async with lock:
                    async with conn.cursor() as cur:
                        try:
                            await cur.execute(self._encode_query(query), values)

                            # Fetch the returned id(s)
                            result = await cur.fetchone()

                            # Check for conflict
                            if result is None:
                                logger.warning(
                                    f"Insert conflict occurred for table '{table}', no ID returned"
                                )
                                raise UniqueConstraintError(
                                    "Insert conflict occurred, no ID returned."
                                )

                            inserted_id = result[0] if result else None

                        except UniqueViolation as e:
                            # Handle unique constraint violation
                            logger.warning(
                                f"Unique constraint violation on table '{table}': {str(e)}"
                            )
                            raise UniqueConstraintError(
                                "Unique constraint violation occurred."
                            )

                        except Exception as e:
                            logger.error(
                                f"Error executing insert on table '{table}': {str(e)}"
                            )
                            raise

                    try:
                        await conn.commit()
                    except Exception as e:
                        logger.error(
                            f"Error committing insert transaction for table '{table}': {str(e)}"
                        )
                        raise

                    return inserted_id

    async def update(
        self,
        table: str,
        set_columns: List[str],
        set_values: List[Any],
        where_clause: str,
        where_params: Optional[Union[Tuple, List]] = None,
    ) -> None:
        """
        Update records in a specified table based on a WHERE clause.

        Args:
            table: The name of the table to update
            set_columns: List of column names to update
            set_values: Values to set for the columns
            where_clause: SQL WHERE clause to identify records to update
            where_params: Parameters for the WHERE clause

        Raises:
            Exception: If update fails
        """
        # Create attributes for span without including potentially sensitive values
        span_attributes = {
            "table": table,
            "column_count": len(set_columns),
            "operation": "UPDATE",
            "has_where_params": where_params is not None,
        }

        with span(
            logger,
            "update_data",
            span_attributes,
            log_entry=False,
            log_exit=False,
        ):
            set_assignments = ", ".join([f"{col} = %s" for col in set_columns])
            query = f"UPDATE {table} SET {set_assignments} WHERE {where_clause}"

            params = set_values + (list(where_params) if where_params else [])
            async with await self.get_connection() as conn:
                # Acquire connection-level lock to prevent concurrent cursor operations
                lock = conn._get_lock()
                async with lock:
                    async with conn.cursor() as cur:
                        try:
                            await cur.execute(query.encode("utf-8"), params)
                        except Exception as e:
                            logger.error(
                                f"Error executing update on table '{table}': {str(e)}"
                            )
                            raise

                    try:
                        await conn.commit()
                    except Exception as e:
                        logger.error(
                            f"Error committing update transaction for table '{table}': {str(e)}"
                        )
                        raise

    async def upsert(
        self,
        table: str,
        columns: List[str],
        values: Union[Any, Tuple[Any, ...]],
        conflict_columns: Optional[List[str]] = None,
        update_columns: Optional[List[str]] = None,
    ) -> Optional[int]:
        """
        Insert data into a specified table or update if it already exists (upsert operation).

        Args:
            table: The name of the table to upsert into
            columns: List of column names
            values: Values to insert/update
            conflict_columns: Columns to check for conflicts (unique constraint)
            update_columns: Columns to update if conflict exists (defaults to all columns except conflict_columns)

        Returns:
            Optional[int]: The ID of the inserted/updated record

        Raises:
            Exception: If upsert operation fails
        """
        # Create attributes for span without including potentially sensitive values
        span_attributes = {
            "table": table,
            "column_count": len(columns),
            "operation": "UPSERT",
            "conflict_column_count": len(conflict_columns) if conflict_columns else 0,
        }

        with span(
            logger,
            "upsert_data",
            span_attributes,
            log_entry=False,
            log_exit=False,
        ):
            # If update_columns not specified, update all columns except conflict columns
            if update_columns is None:
                conflict_column_set = set(conflict_columns or [])
                update_columns = [
                    col
                    for col in columns
                    if col not in conflict_column_set and col != "id"
                ]

            # Create the INSERT part
            placeholders = ", ".join(["%s"] * len(columns))
            columns_formatted = ", ".join(columns)

            # Create the ON CONFLICT part
            conflict_clause = ", ".join(conflict_columns) if conflict_columns else "id"

            # Create the DO UPDATE part
            update_clause = ", ".join(
                [f"{col} = EXCLUDED.{col}" for col in update_columns]
            )

            # Build the complete query
            query = f"""
                INSERT INTO {table} ({columns_formatted}) 
                VALUES ({placeholders})
                ON CONFLICT ({conflict_clause}) 
                DO UPDATE SET {update_clause}
                RETURNING id;
            """

            async with await self.get_connection() as conn:
                # Acquire connection-level lock to prevent concurrent cursor operations
                lock = conn._get_lock()
                async with lock:
                    async with conn.cursor() as cur:
                        try:
                            await cur.execute(self._encode_query(query), values)

                            # Fetch the returned id
                            result = await cur.fetchone()
                            upserted_id = result[0] if result else None
                        except Exception as e:
                            logger.error(
                                f"Error executing upsert on table '{table}': {str(e)}"
                            )
                            raise

                    try:
                        await conn.commit()
                    except Exception as e:
                        logger.error(
                            f"Error committing upsert transaction for table '{table}': {str(e)}"
                        )
                        raise

                return upserted_id

    async def delete(
        self,
        table: str,
        where_clause: str,
        where_params: Optional[Union[Tuple, List]] = None,
    ) -> None:
        """
        Delete records from a specified table based on a WHERE clause.

        Args:
            table: The name of the table to delete from
            where_clause: SQL WHERE clause to identify records to delete
            where_params: Parameters for the WHERE clause

        Raises:
            Exception: If delete operation fails
        """
        # Create attributes for span without including potentially sensitive values
        span_attributes = {
            "table": table,
            "operation": "DELETE",
            "has_where_params": where_params is not None,
        }

        with span(
            logger,
            "delete_data",
            span_attributes,
            log_entry=False,
            log_exit=False,
        ):
            query = f"DELETE FROM {table} WHERE {where_clause}"

            async with await self.get_connection() as conn:
                # Acquire connection-level lock to prevent concurrent cursor operations
                lock = conn._get_lock()
                async with lock:
                    async with conn.cursor() as cur:
                        try:
                            await cur.execute(self._encode_query(query), where_params)
                        except Exception as e:
                            logger.error(
                                f"Error executing delete on table '{table}': {str(e)}"
                            )
                            raise

                    try:
                        await conn.commit()
                    except Exception as e:
                        logger.error(
                            f"Error committing delete transaction for table '{table}': {str(e)}"
                        )
                        raise

    async def close(self) -> None:
        """Close the underlying connection pool if it has been created."""
        with span(
            logger,
            "close_db",
            {"operation": "CLOSE"},
            log_entry=False,
            log_exit=False,
        ):
            while True:
                try:
                    conn = self._pool_queue.get_nowait()
                except asyncio.QueueEmpty:
                    break
                try:
                    await conn.close()
                except Exception:
                    pass
            self._active_connections = 0
            self._connection_timestamps.clear()  # Clean up all timestamps
