# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/utils/psycopg3_optimizations.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

psycopg3-specific optimizations for database operations.

This module provides optimized database operations leveraging psycopg3's
advanced features:
- COPY protocol for bulk inserts (5-10x faster than INSERT)
- Pipeline mode for batch queries (reduced round-trips)
- Prepared statement hints

These optimizations are PostgreSQL-specific and gracefully fall back to
standard SQLAlchemy operations for other databases.

Examples:
    >>> from mcpgateway.utils.psycopg3_optimizations import is_psycopg3_backend
    >>> isinstance(is_psycopg3_backend(), bool)
    True
"""

# Standard
from datetime import datetime
import io
import logging
from typing import Any, Iterable, List, Optional, Sequence, Tuple, TypeVar

# Third-Party
from sqlalchemy import text
from sqlalchemy.orm import Session

logger = logging.getLogger(__name__)

# Check if we're using psycopg3 backend
_is_psycopg3: Optional[bool] = None


def is_psycopg3_backend() -> bool:
    """Check if the current database backend is PostgreSQL with psycopg3.

    Returns:
        True if using PostgreSQL with psycopg3 driver, False otherwise.

    Examples:
        >>> isinstance(is_psycopg3_backend(), bool)
        True
    """
    global _is_psycopg3
    if _is_psycopg3 is None:
        try:
            # First-Party
            from mcpgateway.db import backend, driver

            _is_psycopg3 = backend == "postgresql" and driver in ("psycopg", "default", "")
        except ImportError:
            _is_psycopg3 = False
    return _is_psycopg3


def _format_value_for_copy(value: Any) -> str:
    """Format a Python value for PostgreSQL COPY TEXT format.

    Args:
        value: The value to format.

    Returns:
        String representation suitable for COPY TEXT format.
    """
    if value is None:
        return "\\N"  # NULL representation in COPY
    if isinstance(value, bool):
        return "t" if value else "f"
    if isinstance(value, datetime):
        return value.isoformat()
    if isinstance(value, str):
        # Escape special characters for COPY TEXT format
        return value.replace("\\", "\\\\").replace("\t", "\\t").replace("\n", "\\n").replace("\r", "\\r")
    return str(value)


def bulk_insert_with_copy(
    db: Session,
    table_name: str,
    columns: Sequence[str],
    rows: Iterable[Sequence[Any]],
    schema: Optional[str] = None,
) -> int:
    """Bulk insert rows using PostgreSQL COPY protocol.

    This is significantly faster than individual INSERT statements or even
    bulk_insert_mappings for large datasets. The COPY protocol streams data
    directly to PostgreSQL with minimal overhead.

    Args:
        db: SQLAlchemy session.
        table_name: Name of the target table.
        columns: Sequence of column names to insert.
        rows: Iterable of row tuples matching column order.
        schema: Optional schema name (defaults to search_path).

    Returns:
        Number of rows inserted.

    Note:
        Falls back to executemany for non-PostgreSQL databases.
    """
    if not is_psycopg3_backend():
        # Fallback to standard INSERT for non-PostgreSQL
        return _bulk_insert_fallback(db, table_name, columns, rows, schema)

    try:
        # Get raw psycopg connection from SQLAlchemy
        raw_conn = db.connection().connection.dbapi_connection

        # Build the qualified table name
        qualified_table = f"{schema}.{table_name}" if schema else table_name
        columns_str = ", ".join(columns)

        # Create a file-like object with COPY data
        buffer = io.StringIO()
        row_count = 0

        for row in rows:
            line = "\t".join(_format_value_for_copy(v) for v in row)
            buffer.write(line + "\n")
            row_count += 1

        if row_count == 0:
            return 0

        buffer.seek(0)

        # Use psycopg3's COPY FROM
        with raw_conn.cursor() as cur:
            with cur.copy(f"COPY {qualified_table} ({columns_str}) FROM STDIN") as copy:
                while data := buffer.read(8192):
                    copy.write(data)

        logger.debug("COPY inserted %d rows into %s", row_count, qualified_table)
        return row_count

    except Exception as e:
        logger.warning("COPY failed, falling back to INSERT: %s", e)
        return _bulk_insert_fallback(db, table_name, columns, rows, schema)


def _bulk_insert_fallback(
    db: Session,
    table_name: str,
    columns: Sequence[str],
    rows: Iterable[Sequence[Any]],
    schema: Optional[str] = None,
) -> int:
    """Fallback bulk insert using executemany.

    Args:
        db: SQLAlchemy session.
        table_name: Name of the target table.
        columns: Sequence of column names to insert.
        rows: Iterable of row tuples matching column order.
        schema: Optional schema name.

    Returns:
        Number of rows inserted.
    """
    qualified_table = f"{schema}.{table_name}" if schema else table_name
    columns_str = ", ".join(columns)
    placeholders = ", ".join(f":{col}" for col in columns)

    sql = text(f"INSERT INTO {qualified_table} ({columns_str}) VALUES ({placeholders})")  # nosec B608 - table/columns from SQLAlchemy models, not user input

    row_list = list(rows)
    if not row_list:
        return 0

    # Convert rows to list of dicts
    data = [dict(zip(columns, row)) for row in row_list]
    db.execute(sql, data)

    return len(row_list)


T = TypeVar("T")


def execute_pipelined(
    db: Session,
    queries: Sequence[Tuple[str, dict]],
) -> List[List[Any]]:
    """Execute multiple queries in pipeline mode for reduced round-trips.

    Pipeline mode allows sending multiple queries without waiting for
    individual responses, significantly reducing latency for independent
    queries.

    Args:
        db: SQLAlchemy session.
        queries: Sequence of (sql_string, params_dict) tuples.

    Returns:
        List of result lists, one per query.

    Note:
        Falls back to sequential execution for non-PostgreSQL databases.
    """
    if not queries:
        return []

    if not is_psycopg3_backend():
        # Fallback to sequential execution
        return [list(db.execute(text(sql), params).fetchall()) for sql, params in queries]

    try:
        raw_conn = db.connection().connection.dbapi_connection

        results = []
        with raw_conn.pipeline():
            cursors = []
            for sql, params in queries:
                cur = raw_conn.execute(sql, params)
                cursors.append(cur)

            for cur in cursors:
                try:
                    results.append(list(cur.fetchall()))
                except Exception:
                    results.append([])

        logger.debug("Pipelined %d queries", len(queries))
        return results

    except Exception as e:
        logger.warning("Pipeline mode failed, falling back to sequential: %s", e)
        return [list(db.execute(text(sql), params).fetchall()) for sql, params in queries]


def bulk_insert_metrics(
    db: Session,
    table_name: str,
    metrics: Sequence[dict],
    columns: Optional[Sequence[str]] = None,
) -> int:
    """Optimized bulk insert for metric records.

    Uses COPY protocol on PostgreSQL for maximum performance when
    writing metrics data.

    Args:
        db: SQLAlchemy session.
        table_name: Name of the metrics table.
        metrics: Sequence of metric dictionaries.
        columns: Optional explicit column list. If not provided,
                 uses keys from first metric dict.

    Returns:
        Number of rows inserted.

    Examples:
        >>> # Example usage (would need actual DB connection):
        >>> # bulk_insert_metrics(db, "tool_metrics", [
        >>> #     {"tool_id": "abc", "timestamp": datetime.now(), "response_time": 0.5, "is_success": True}
        >>> # ])
    """
    if not metrics:
        return 0

    # Determine columns from first metric if not provided
    if columns is None:
        columns = list(metrics[0].keys())

    # Convert dicts to ordered tuples
    rows = [[m.get(col) for col in columns] for m in metrics]

    return bulk_insert_with_copy(db, table_name, columns, rows)


def get_raw_connection(db: Session) -> Any:
    """Get the raw psycopg3 connection from a SQLAlchemy session.

    Args:
        db: SQLAlchemy session.

    Returns:
        The underlying psycopg3 connection, or None if not using psycopg3.
    """
    if not is_psycopg3_backend():
        return None

    try:
        return db.connection().connection.dbapi_connection
    except Exception:
        return None
