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

Database query logging middleware for N+1 detection.
This middleware logs all database queries per request to help identify
N+1 query patterns and other performance issues.

Enable with:
    DB_QUERY_LOG_ENABLED=true

Output files:
    - logs/db-queries.log (human-readable text)
    - logs/db-queries.jsonl (JSON Lines for tooling)
"""

# Standard
from contextvars import ContextVar
from datetime import datetime, timezone
import logging
from pathlib import Path
import re
import threading
import time
from typing import Any, Dict, List, Optional, Pattern

# Third-Party
import orjson
from sqlalchemy import event
from sqlalchemy.engine import Engine
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response

# First-Party
from mcpgateway.config import get_settings
from mcpgateway.middleware.path_filter import should_skip_db_query_logging

logger = logging.getLogger(__name__)

# ============================================================================
# Precompiled regex patterns for query normalization (compiled once at module load)
# ============================================================================
_QUOTED_STRING_RE: Pattern[str] = re.compile(r"'[^']*'")
_NUMBER_RE: Pattern[str] = re.compile(r"\b\d+\b")
_IN_CLAUSE_RE: Pattern[str] = re.compile(r"IN\s*\([^)]+\)", re.IGNORECASE)
_WHITESPACE_RE: Pattern[str] = re.compile(r"\s+")
_TABLE_NAME_RE: Pattern[str] = re.compile(r"(?:FROM|INTO|UPDATE)\s+[\"']?(\w+)[\"']?", re.IGNORECASE)

# Context variable to track queries per request
_request_context: ContextVar[Optional[Dict[str, Any]]] = ContextVar("db_query_request_context", default=None)

# Lock for thread-safe file writing
_file_lock = threading.Lock()

# Track if we've already instrumented the engine
_instrumented_engines: set = set()


def _normalize_query(sql: str) -> str:
    """Normalize a SQL query for pattern detection.

    Replaces specific values with placeholders to identify similar queries.
    Uses precompiled regex patterns for performance.

    Args:
        sql: The SQL query string

    Returns:
        Normalized query string
    """
    # Replace quoted strings (uses precompiled regex)
    normalized = _QUOTED_STRING_RE.sub("'?'", sql)
    # Replace numbers (uses precompiled regex)
    normalized = _NUMBER_RE.sub("?", normalized)
    # Replace IN clauses with multiple values (uses precompiled regex)
    normalized = _IN_CLAUSE_RE.sub("IN (?)", normalized)
    # Normalize whitespace (uses precompiled regex)
    normalized = _WHITESPACE_RE.sub(" ", normalized).strip()
    return normalized


def _extract_table_name(sql: str) -> Optional[str]:
    """Extract the main table name from a SQL query.

    Uses precompiled regex pattern for performance.

    Args:
        sql: The SQL query string

    Returns:
        Table name or None
    """
    # Match FROM table or INTO table or UPDATE table (uses precompiled regex)
    match = _TABLE_NAME_RE.search(sql)
    if match:
        return match.group(1)
    return None


def _detect_n1_patterns(queries: List[Dict[str, Any]], threshold: int = 3) -> List[Dict[str, Any]]:
    """Detect potential N+1 query patterns.

    Args:
        queries: List of query dictionaries with 'sql' key
        threshold: Minimum repetitions to flag as N+1

    Returns:
        List of detected N+1 patterns with details
    """
    patterns: Dict[str, List[int]] = {}

    for idx, q in enumerate(queries):
        normalized = _normalize_query(q.get("sql", ""))
        if normalized not in patterns:
            patterns[normalized] = []
        patterns[normalized].append(idx)

    n1_issues = []
    for pattern, indices in patterns.items():
        if len(indices) >= threshold:
            table = _extract_table_name(pattern)
            n1_issues.append(
                {
                    "pattern": pattern[:200],  # Truncate long patterns
                    "count": len(indices),
                    "table": table,
                    "query_indices": indices,
                }
            )

    return sorted(n1_issues, key=lambda x: x["count"], reverse=True)


def _format_text_log(request_data: Dict[str, Any], queries: List[Dict[str, Any]], n1_issues: List[Dict[str, Any]]) -> str:
    """Format request and queries as human-readable text.

    Args:
        request_data: Request metadata
        queries: List of executed queries
        n1_issues: Detected N+1 patterns

    Returns:
        Formatted text string
    """
    lines = []
    separator = "=" * 80

    # Header
    lines.append(separator)
    timestamp = request_data.get("timestamp", datetime.now(timezone.utc).isoformat())
    method = request_data.get("method", "?")
    path = request_data.get("path", "?")
    lines.append(f"[{timestamp}] {method} {path}")

    # Metadata line
    meta_parts = []
    if request_data.get("user"):
        meta_parts.append(f"User: {request_data['user']}")
    if request_data.get("correlation_id"):
        meta_parts.append(f"Correlation-ID: {request_data['correlation_id']}")
    meta_parts.append(f"Queries: {len(queries)}")
    total_ms = sum(q.get("duration_ms", 0) for q in queries)
    meta_parts.append(f"Total: {total_ms:.1f}ms")
    lines.append(" | ".join(meta_parts))
    lines.append(separator)

    # N+1 warnings at top if detected
    if n1_issues:
        lines.append("")
        lines.append("⚠️  POTENTIAL N+1 QUERIES DETECTED:")
        for issue in n1_issues:
            table_info = f" on '{issue['table']}'" if issue.get("table") else ""
            lines.append(f"   • {issue['count']}x similar queries{table_info}")
            lines.append(f"     Pattern: {issue['pattern'][:100]}...")
        lines.append("")

    # Query list
    for idx, q in enumerate(queries, 1):
        duration = q.get("duration_ms", 0)
        sql = q.get("sql", "")

        # Check if this query is part of an N+1 pattern
        n1_marker = ""
        for issue in n1_issues:
            if idx - 1 in issue.get("query_indices", []):
                n1_marker = "  ← N+1"
                break

        # Truncate long queries
        if len(sql) > 200:
            sql = sql[:200] + "..."

        lines.append(f"  {idx:3}. [{duration:6.1f}ms] {sql}{n1_marker}")

    # Footer
    lines.append("-" * 80)
    if n1_issues:
        lines.append(f"⚠️  {len(n1_issues)} potential N+1 pattern(s) detected - see docs/docs/development/db-performance.md")
    lines.append(f"Total: {len(queries)} queries, {total_ms:.1f}ms")
    lines.append(separator)
    lines.append("")

    return "\n".join(lines)


def _format_json_log(request_data: Dict[str, Any], queries: List[Dict[str, Any]], n1_issues: List[Dict[str, Any]]) -> str:
    """Format request and queries as JSON.

    Args:
        request_data: Request metadata
        queries: List of executed queries
        n1_issues: Detected N+1 patterns

    Returns:
        JSON string (single line)
    """
    total_ms = sum(q.get("duration_ms", 0) for q in queries)

    log_entry = {
        "timestamp": request_data.get("timestamp", datetime.now(timezone.utc).isoformat()),
        "method": request_data.get("method"),
        "path": request_data.get("path"),
        "user": request_data.get("user"),
        "correlation_id": request_data.get("correlation_id"),
        "status_code": request_data.get("status_code"),
        "query_count": len(queries),
        "total_query_ms": round(total_ms, 2),
        "request_duration_ms": request_data.get("request_duration_ms"),
        "n1_issues": n1_issues if n1_issues else None,
        "queries": [
            {
                "sql": q.get("sql", "")[:500],  # Truncate long queries
                "duration_ms": round(q.get("duration_ms", 0), 2),
                "table": _extract_table_name(q.get("sql", "")),
            }
            for q in queries
        ],
    }

    return orjson.dumps(log_entry, default=str).decode()


def _write_logs(request_data: Dict[str, Any], queries: List[Dict[str, Any]]) -> None:
    """Write query logs to file(s).

    Args:
        request_data: Request metadata
        queries: List of executed queries
    """
    settings = get_settings()

    # Skip if no queries or below threshold
    if not queries or len(queries) < settings.db_query_log_min_queries:
        return

    # Detect N+1 patterns
    n1_issues = []
    if settings.db_query_log_detect_n1:
        n1_issues = _detect_n1_patterns(queries, settings.db_query_log_n1_threshold)

    log_format = settings.db_query_log_format.lower()

    with _file_lock:
        # Write text log
        if log_format in ("text", "both"):
            text_path = Path(settings.db_query_log_file)
            text_path.parent.mkdir(parents=True, exist_ok=True)
            with open(text_path, "a", encoding="utf-8") as f:
                f.write(_format_text_log(request_data, queries, n1_issues))

        # Write JSON log
        if log_format in ("json", "both"):
            json_path = Path(settings.db_query_log_json_file)
            json_path.parent.mkdir(parents=True, exist_ok=True)
            with open(json_path, "a", encoding="utf-8") as f:
                f.write(_format_json_log(request_data, queries, n1_issues) + "\n")


def _before_cursor_execute(conn: Any, _cursor: Any, _statement: str, _parameters: Any, _context: Any, _executemany: bool) -> None:
    """SQLAlchemy event handler called before query execution.

    Args:
        conn: Database connection
        _cursor: Database cursor (unused, required by SQLAlchemy event signature)
        _statement: SQL statement to execute (unused, required by SQLAlchemy event signature)
        _parameters: Query parameters (unused, required by SQLAlchemy event signature)
        _context: Execution context (unused, required by SQLAlchemy event signature)
        _executemany: Whether this is an executemany call (unused, required by SQLAlchemy event signature)
    """
    ctx = _request_context.get()
    if ctx is None:
        return

    # Store start time on the connection
    conn.info["_query_start_time"] = time.perf_counter()


# Tables to exclude from query logging (internal/observability tables)
_EXCLUDED_TABLES = {
    "observability_traces",
    "observability_spans",
    "observability_events",
    "observability_metrics",
    "structured_log_entries",
    "audit_logs",
    "security_events",
}


def _should_exclude_query(statement: str) -> bool:
    """Check if a query should be excluded from logging.

    Args:
        statement: SQL statement

    Returns:
        True if the query should be excluded
    """
    statement_upper = statement.upper()
    for table in _EXCLUDED_TABLES:
        if table.upper() in statement_upper:
            return True
    return False


def _after_cursor_execute(conn: Any, _cursor: Any, statement: str, parameters: Any, _context: Any, executemany: bool) -> None:
    """SQLAlchemy event handler called after query execution.

    Args:
        conn: Database connection
        _cursor: Database cursor (unused, required by SQLAlchemy event signature)
        statement: SQL statement that was executed
        parameters: Query parameters
        _context: Execution context (unused, required by SQLAlchemy event signature)
        executemany: Whether this was an executemany call
    """
    ctx = _request_context.get()
    if ctx is None:
        return

    # Skip internal observability queries
    if _should_exclude_query(statement):
        conn.info.pop("_query_start_time", None)  # Clean up
        return

    # Calculate duration
    start_time = conn.info.pop("_query_start_time", None)
    duration_ms = (time.perf_counter() - start_time) * 1000 if start_time else 0

    # Get settings for parameter inclusion
    settings = get_settings()

    query_info = {
        "sql": statement,
        "duration_ms": duration_ms,
        "executemany": executemany,
    }

    if settings.db_query_log_include_params and parameters:
        # Sanitize parameters - don't include actual values by default
        query_info["param_count"] = len(parameters) if isinstance(parameters, (list, tuple, dict)) else 1

    ctx["queries"].append(query_info)


def instrument_engine_for_logging(engine: Engine) -> None:
    """Instrument a SQLAlchemy engine for query logging.

    Args:
        engine: SQLAlchemy engine to instrument
    """
    engine_id = id(engine)
    if engine_id in _instrumented_engines:
        return

    event.listen(engine, "before_cursor_execute", _before_cursor_execute)
    event.listen(engine, "after_cursor_execute", _after_cursor_execute)
    _instrumented_engines.add(engine_id)
    logger.info("Database query logging instrumentation enabled")


class DBQueryLoggingMiddleware(BaseHTTPMiddleware):
    """Middleware to log database queries per request.

    This middleware:
    1. Creates a request context to collect queries
    2. Captures request metadata (method, path, user, correlation ID)
    3. After the request, writes all queries to log file(s)
    4. Detects and flags potential N+1 query patterns
    """

    async def dispatch(self, request: Request, call_next) -> Response:
        """Process request and log database queries.

        Args:
            request: The incoming request
            call_next: Next middleware/handler

        Returns:
            Response from the handler
        """
        settings = get_settings()

        if not settings.db_query_log_enabled:
            return await call_next(request)

        # Skip static files and health checks
        path = request.url.path
        if should_skip_db_query_logging(path):
            return await call_next(request)

        # Create request context
        ctx: Dict[str, Any] = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "method": request.method,
            "path": path,
            "user": None,
            "correlation_id": request.headers.get(settings.correlation_id_header),
            "queries": [],
        }

        # Try to get user from request state (set by auth middleware)
        if hasattr(request.state, "user"):
            ctx["user"] = getattr(request.state.user, "username", str(request.state.user))
        elif hasattr(request.state, "username"):
            ctx["user"] = request.state.username

        # Set context for SQLAlchemy event handlers
        token = _request_context.set(ctx)

        try:
            start_time = time.perf_counter()
            response = await call_next(request)
            request_duration = (time.perf_counter() - start_time) * 1000

            ctx["status_code"] = response.status_code
            ctx["request_duration_ms"] = round(request_duration, 2)

            return response
        finally:
            # Write logs
            try:
                _write_logs(ctx, ctx["queries"])
            except Exception as e:
                logger.warning(f"Failed to write query log: {e}")

            # Reset context
            _request_context.reset(token)


def setup_query_logging(app: Any, engine: Engine) -> None:
    """Set up database query logging for an application.

    Args:
        app: FastAPI application
        engine: SQLAlchemy engine
    """
    settings = get_settings()

    if not settings.db_query_log_enabled:
        return

    # Instrument the engine
    instrument_engine_for_logging(engine)

    # Add middleware
    app.add_middleware(DBQueryLoggingMiddleware)

    logger.info(f"Database query logging enabled: " f"format={settings.db_query_log_format}, " f"text_file={settings.db_query_log_file}, " f"json_file={settings.db_query_log_json_file}")
