from flask import Flask, request, jsonify, make_response
from prometheus_api_client import PrometheusConnect
import csv
import io
from datetime import datetime, timezone, timedelta
import logging
import os
import hmac
import re
import threading
import time
import boto3
from requests_aws4auth import AWS4Auth
import psycopg2
import psycopg2.extras

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# Allowlist for the /execute-query debug endpoint: only read-only statement types.
_SQL_ALLOWLIST_RE = re.compile(r'^\s*(SELECT|EXPLAIN|SHOW)\b', re.IGNORECASE)

# Secondary blocklist: reject sensitive system functions/views and WITH clauses.
# pg_read_file/pg_shadow/pg_authid/pg_read_binary_file/pg_ls_dir/pg_stat_file can
# expose credentials or filesystem contents; WITH blocks writable CTEs
# (INSERT/UPDATE/DELETE inside a CTE that passes the SELECT allowlist via
# "WITH ... SELECT").
_SQL_BLOCKLIST_RE = re.compile(
    r'\b(pg_read_file|pg_read_binary_file|pg_ls_dir|pg_stat_file|pg_shadow|pg_authid|dblink|lo_export|lo_import|COPY|DO|CALL)\b|\bWITH\b',
    re.IGNORECASE,
)

# Regex for quote-aware SQL comment stripping.  Matches (in priority order):
#   1. Single-quoted string literals (including escaped quotes '')  -> kept
#   2. Block comments /* ... */                                     -> removed
#   3. Single-line comments -- ...                                  -> removed
_SQL_COMMENT_RE = re.compile(
    r"'[^']*(?:''[^']*)*'"   # single-quoted literal (handles '' escapes)
    r"|/\*.*?\*/"             # block comment
    r"|--[^\n]*",             # line comment
    re.DOTALL,
)


def _strip_sql_comments(sql: str) -> str:
    """Remove SQL comments while preserving single-quoted string literals.

    Naive regex-based comment stripping (e.g. ``re.sub(r'--[^\\n]*', ...)``)
    can be tricked by placing comment delimiters inside string literals:

        SELECT 'x' || '--' || (SELECT * FROM pg_shadow)

    This function matches quoted strings first (and keeps them intact) so that
    only *real* comments are removed.
    """
    def _replace(m: re.Match) -> str:
        text = m.group()
        if text.startswith("'"):
            return text          # keep string literals
        return ""                # remove comments

    return _SQL_COMMENT_RE.sub(_replace, sql)


try:
    from monitoring_flask_backend.promql_utils import escape_promql_label, escape_promql_regex_literal  # noqa: F401
except ModuleNotFoundError:
    # Running directly from the monitoring_flask_backend/ directory (e.g. pytest from that dir)
    from promql_utils import escape_promql_label, escape_promql_regex_literal  # noqa: F401


def smart_truncate_query(query: str, max_length: int = 40) -> str:
    """
    Smart SQL query truncation for display names.

    Provides more informative truncation than simple character limits by:
    1. Stripping leading comments (/* ... */ and -- ...)
    2. For SELECT queries: showing "SELECT ... FROM <table_names>"
    3. For CTEs: showing "WITH cte1, cte2 SELECT ... FROM ..."
    4. Fallback to simple truncation on any parse error

    Args:
        query: The SQL query text
        max_length: Maximum length for the result (default: 40)

    Returns:
        Truncated query string suitable for display
    """
    if not query:
        return ''

    # Store original for fallback (before try block to ensure it's always bound)
    original_query = query

    try:
        # Step 1: Strip ALL block comments /* ... */ (not just leading ones)
        # This handles inline comments like: SELECT /* comment */ FROM ...
        query = re.sub(r'/\*.*?\*/', ' ', query, flags=re.DOTALL)

        # Step 2: Strip ALL single-line comments -- ... to end of line
        query = re.sub(r'--[^\n]*', ' ', query)

        # Normalize whitespace
        query = re.sub(r'\s+', ' ', query).strip()

        # If already short enough, return it
        if len(query) <= max_length:
            return query

        # Step 3: Parse CTEs (WITH clause)
        cte_names = []
        main_query = query  # The main SELECT statement (after CTEs)
        cte_match = re.match(r'^WITH\s+(RECURSIVE\s+)?', query, re.IGNORECASE)
        if cte_match:
            # Find the main SELECT after CTEs by counting parentheses
            # CTEs are: WITH name AS (...), name2 AS (...) SELECT ...
            after_with = query[cte_match.end():]

            # Extract CTE names from the section before the main SELECT
            cte_name_matches = re.findall(r'(\w+)\s+AS\s*\(', after_with, re.IGNORECASE)
            cte_names = cte_name_matches

            # Find the main SELECT by tracking parentheses depth
            # The main SELECT is at depth 0 after all CTE definitions
            paren_depth = 0
            main_select_pos = None
            i = 0
            while i < len(after_with):
                if after_with[i] == '(':
                    paren_depth += 1
                elif after_with[i] == ')':
                    paren_depth -= 1
                elif paren_depth == 0:
                    # Check if we're at a SELECT keyword at depth 0
                    remaining = after_with[i:].upper()
                    if remaining.startswith('SELECT'):
                        main_select_pos = i
                        break
                i += 1

            if main_select_pos is not None:
                main_query = after_with[main_select_pos:]

        # Step 4: Extract FROM clause tables from the main query
        # Use a more robust approach: find FROM and extract until a SQL keyword
        from_tables = []
        from_match = re.search(r'\bFROM\s+', main_query, re.IGNORECASE)
        if from_match:
            # Get everything after FROM
            after_from = main_query[from_match.end():]
            # Find where the FROM clause ends (at a SQL keyword or end of string)
            end_match = re.search(
                r'\b(WHERE|JOIN|LEFT|RIGHT|INNER|OUTER|CROSS|FULL|NATURAL|ORDER|GROUP|HAVING|LIMIT|OFFSET|UNION|INTERSECT|EXCEPT|FOR|RETURNING)\b',
                after_from, re.IGNORECASE
            )
            if end_match:
                from_clause = after_from[:end_match.start()].strip()
            else:
                from_clause = after_from.strip()

            # Extract table names (handle aliases, schemas, commas)
            # Split by comma and extract first word of each (the table/alias)
            table_parts = from_clause.split(',')
            for part in table_parts:
                part = part.strip()
                if part:
                    # Get the table name (first word, might include schema.table)
                    table_name_match = re.match(r'^([\w.]+)', part)
                    if table_name_match:
                        table_name = table_name_match.group(1)
                        # Skip subqueries (start with parentheses)
                        if not part.startswith('('):
                            from_tables.append(table_name)

        # Step 5: Determine query type and build display string
        query_upper = query.upper()

        if query_upper.startswith('SELECT') or cte_names:
            # Build the smart truncated version (lowercase to match pg_stat_statements)
            if cte_names and from_tables:
                result = f"with {', '.join(cte_names)} select ... from {', '.join(from_tables)}"
            elif cte_names:
                result = f"with {', '.join(cte_names)} select ..."
            elif from_tables:
                result = f"select ... from {', '.join(from_tables)}"
            else:
                # No FROM clause found, fall back to simple truncation
                result = query[:max_length - 3] + '...'

            # If result is still too long, truncate it
            if len(result) > max_length:
                result = result[:max_length - 3] + '...'

            return result

        elif query_upper.startswith('INSERT'):
            # For INSERT, show the table name (lowercase to match pg_stat_statements)
            insert_match = re.match(r'^INSERT\s+INTO\s+(\S+)', query, re.IGNORECASE)
            if insert_match:
                result = f"insert into {insert_match.group(1)} ..."
                if len(result) > max_length:
                    result = result[:max_length - 3] + '...'
                return result

        elif query_upper.startswith('UPDATE'):
            # For UPDATE, show the table name (lowercase to match pg_stat_statements)
            update_match = re.match(r'^UPDATE\s+(\S+)', query, re.IGNORECASE)
            if update_match:
                result = f"update {update_match.group(1)} ..."
                if len(result) > max_length:
                    result = result[:max_length - 3] + '...'
                return result

        elif query_upper.startswith('DELETE'):
            # For DELETE, show the table name (lowercase to match pg_stat_statements)
            delete_match = re.match(r'^DELETE\s+FROM\s+(\S+)', query, re.IGNORECASE)
            if delete_match:
                result = f"delete from {delete_match.group(1)} ..."
                if len(result) > max_length:
                    result = result[:max_length - 3] + '...'
                return result

        # Fallback: simple truncation
        if len(query) > max_length:
            return query[:max_length - 3] + '...'
        return query

    except Exception:
        # On any error, fall back to simple truncation of original
        if len(original_query) > max_length:
            return original_query[:max_length - 3] + '...'
        return original_query


app = Flask(__name__)

# PostgreSQL sink connection for query text lookups
POSTGRES_SINK_URL = os.environ.get('POSTGRES_SINK_URL', 'postgresql://pgwatch@sink-postgres:5432/measurements')

# PostgreSQL direct connection config (used by debug endpoints only)
DB_CONFIG = {
    'host': os.environ.get('DB_HOST', 'localhost'),
    'dbname': os.environ.get('DB_NAME', 'postgres'),
    'user': os.environ.get('DB_USER'),
    'password': os.environ.get('DB_PASSWORD'),
}

# Secret key required to access debug endpoints (must be set to enable debug auth)
_DEBUG_SECRET_KEY = os.environ.get('DEBUG_SECRET_KEY', '')

# Prometheus connection - use environment variable with fallback
PROMETHEUS_URL = os.environ.get('PROMETHEUS_URL', 'http://localhost:8428')

# Retention: how long to keep stale queryid entries (queryids no longer in pg_stat_statements).
# Default 30 days. Rows older than this are deleted by background cleanup.
QUERYID_RETENTION_HOURS = max(1, int(os.environ.get('QUERYID_RETENTION_HOURS', '720')))

# How recent a queryid must be to appear in /query_info_metrics (must be > pgwatch collection interval).
# Default 10 minutes. Queryids not seen within this window are excluded from metrics export.
QUERYID_ACTIVE_MINUTES = max(1, int(os.environ.get('QUERYID_ACTIVE_MINUTES', '10')))

# Rows deleted per batch during retention cleanup. Limits lock contention on large tables.
QUERYID_RETENTION_BATCH_SIZE = max(1, int(os.environ.get('QUERYID_RETENTION_BATCH_SIZE', '10000')))


# Per-worker flag; each gunicorn worker re-runs the (cheap, idempotent)
# verification once on first request after startup. Cross-worker
# serialization is the advisory lock below — do not try to replace this
# with a cross-process cache.
_trigger_migration_applied = False
_trigger_migration_lock = threading.Lock()
_trigger_migration_warning_state = {'last_at': 0.0}
_TRIGGER_MIGRATION_WARNING_INTERVAL_SECONDS = 300

# Arbitrary constant keys for pg_advisory_lock — serialize across workers
# so only one gunicorn process at a time runs the DDL / retention cleanup.
# Different keys so the two paths don't block each other.
_TRIGGER_MIGRATION_ADVISORY_LOCK_KEY = 0x70676169_71756964  # "pgai" + "quid"
_RETENTION_CLEANUP_ADVISORY_LOCK_KEY = 0x70676169_72746e63  # "pgai" + "rtnc"


def _warn_trigger_migration_failure(message: str, *args) -> None:
    """Rate-limit retryable trigger verification warnings per worker."""
    now = time.monotonic()
    if now - _trigger_migration_warning_state['last_at'] < _TRIGGER_MIGRATION_WARNING_INTERVAL_SECONDS:
        return
    _trigger_migration_warning_state['last_at'] = now
    logger.warning(message, *args)


def _apply_trigger_migration() -> None:
    """
    Verify the sink-side dedup trigger matches what this Flask backend
    expects. The function body + trigger are created by
    config/sink-postgres/init.sql as the bootstrap superuser; this call
    only checks that the deployed function contains the advisory-lock
    deduplication path and is named as expected.

    Running as the pgwatch role (which does not own the function) means we
    cannot self-heal an out-of-date deployment — by design. If the
    verification fails, the operator must re-run init.sql as the
    bootstrap role.

    Called lazily on first request. Concurrency-safe: in-process Lock
    for thread-level serialization, Postgres advisory lock for
    worker-level serialization.
    """
    global _trigger_migration_applied
    # Fast path — avoid taking the lock once verification has succeeded.
    if _trigger_migration_applied:
        return

    with _trigger_migration_lock:
        if _trigger_migration_applied:
            return

        conn = None
        try:
            conn = psycopg2.connect(POSTGRES_SINK_URL, connect_timeout=10)
            conn.autocommit = True
            with conn.cursor() as cursor:
                cursor.execute(
                    "select pg_try_advisory_lock(%s)",
                    (_TRIGGER_MIGRATION_ADVISORY_LOCK_KEY,),
                )
                got_lock = cursor.fetchone()[0]
                if not got_lock:
                    logger.debug(
                        "Trigger verification advisory lock held by another worker; skipping"
                    )
                    return

                try:
                    # Verify the function body contains the advisory-lock
                    # delete-or-insert path compatible with pgwatch's
                    # dbname/time partition hierarchy.
                    cursor.execute("""
                        select prosrc
                        from pg_proc
                        where proname = 'enforce_queryid_uniqueness'
                          and pronamespace = 'public'::regnamespace
                    """)
                    row = cursor.fetchone()
                    if row is None:
                        _warn_trigger_migration_failure(
                            "Sink dedup function public.enforce_queryid_uniqueness not found; "
                            "re-run config/sink-postgres/init.sql as the bootstrap role."
                        )
                        return
                    function_body = row[0] or ''
                    function_body_lower = function_body.lower()
                    has_advisory_lock = 'pg_advisory_xact_lock' in function_body_lower
                    has_queryid_lookup = "data->>'queryid'" in function_body_lower
                    has_delete_dedup = re.search(
                        r'\bdelete\s+from\s+(?:public\.)?pgss_queryid_queries\b',
                        function_body,
                        re.IGNORECASE,
                    )
                    if not (has_advisory_lock and has_queryid_lookup and has_delete_dedup):
                        _warn_trigger_migration_failure(
                            "Sink dedup function public.enforce_queryid_uniqueness is missing "
                            "the partition-safe advisory-lock dedup path. Re-run "
                            "config/sink-postgres/init.sql as the bootstrap role to upgrade."
                        )
                        return

                    # Verify the trigger exists with the expected name.
                    cursor.execute("""
                        select 1 from pg_trigger
                        where tgname = 'enforce_queryid_uniqueness_trigger'
                          and tgrelid = 'public.pgss_queryid_queries'::regclass
                    """)
                    if cursor.fetchone() is None:
                        _warn_trigger_migration_failure(
                            "Sink dedup trigger enforce_queryid_uniqueness_trigger is missing "
                            "on public.pgss_queryid_queries. Re-run "
                            "config/sink-postgres/init.sql as the bootstrap role."
                        )
                        return
                finally:
                    cursor.execute(
                        "select pg_advisory_unlock(%s)",
                        (_TRIGGER_MIGRATION_ADVISORY_LOCK_KEY,),
                    )
            _trigger_migration_applied = True
            logger.info("Verified sink dedup function + trigger")
        except Exception as e:
            logger.warning("Failed to verify sink dedup migration: %s", e)
        finally:
            if conn:
                conn.close()


@app.before_request
def _ensure_trigger_migration() -> None:
    """Verify the sink dedup migration lazily on first request."""
    _apply_trigger_migration()

# Guard to prevent unbounded background thread creation per scrape.
# The lock makes the Event check/set an atomic in-process compare-and-set.
_cleanup_running = threading.Event()
_cleanup_start_lock = threading.Lock()


# Cap iterations per run so a huge backlog can't monopolize a connection.
# At BATCH_SIZE=10 000 and 10 iterations this deletes up to 100 000 rows
# per scrape — enough to keep up with typical write rates without holding
# locks too long. Clamped to >= 1 so a misconfiguration doesn't silently
# disable cleanup.
QUERYID_RETENTION_MAX_ITERATIONS = max(
    1, int(os.environ.get('QUERYID_RETENTION_MAX_ITERATIONS', '10'))
)


def _run_retention_cleanup() -> None:
    """
    Delete stale queryid entries in a background thread using a separate
    connection. This prevents the retention DELETE from adding latency to
    every Prometheus scrape of /query_info_metrics.

    Loops up to QUERYID_RETENTION_MAX_ITERATIONS so cleanup can catch up
    on a backlog (e.g. after an outage). statement_timeout/lock_timeout
    prevent the daemon thread from hanging indefinitely on a lock wait.

    Concurrency: _cleanup_running coalesces threads inside one worker;
    pg_try_advisory_lock serializes across workers so only one of N
    gunicorn processes runs the DELETE loop at any moment.
    """
    conn = None
    try:
        conn = psycopg2.connect(
            POSTGRES_SINK_URL,
            connect_timeout=10,
            # Bound each DELETE at the server side so a blocked lock or
            # runaway plan doesn't leave the cleanup thread stuck forever.
            options='-c statement_timeout=60000 -c lock_timeout=5000',
        )
        conn.autocommit = True
        with conn.cursor() as cursor:
            cursor.execute(
                "select pg_try_advisory_lock(%s)",
                (_RETENTION_CLEANUP_ADVISORY_LOCK_KEY,),
            )
            got_lock = cursor.fetchone()[0]
            if not got_lock:
                logger.debug(
                    "Retention cleanup advisory lock held by another worker; skipping"
                )
                return

            total_deleted = 0
            try:
                for _ in range(QUERYID_RETENTION_MAX_ITERATIONS):
                    cursor.execute("""
                        delete from public.pgss_queryid_queries
                        where ctid in (
                            select ctid from public.pgss_queryid_queries
                            where time < now() - make_interval(hours => %s)
                            limit %s
                        )
                    """, (QUERYID_RETENTION_HOURS, QUERYID_RETENTION_BATCH_SIZE))
                    deleted = cursor.rowcount
                    total_deleted += deleted
                    if deleted < QUERYID_RETENTION_BATCH_SIZE:
                        break
            finally:
                cursor.execute(
                    "select pg_advisory_unlock(%s)",
                    (_RETENTION_CLEANUP_ADVISORY_LOCK_KEY,),
                )
            if total_deleted > 0:
                logger.info(
                    "Retention cleanup: deleted %d stale queryid entries",
                    total_deleted,
                )
    except Exception as e:
        logger.warning("Retention cleanup failed: %s", e)
    finally:
        if conn:
            conn.close()


def _run_retention_cleanup_guarded() -> None:
    """Wrapper that clears the _cleanup_running flag when done."""
    try:
        _run_retention_cleanup()
    finally:
        _cleanup_running.clear()


def _start_retention_cleanup_thread() -> None:
    """Start one in-process retention cleanup thread if none is running."""
    with _cleanup_start_lock:
        if _cleanup_running.is_set():
            return

        _cleanup_running.set()
        try:
            threading.Thread(
                target=_run_retention_cleanup_guarded,
                name="retention-cleanup",
                daemon=True,
            ).start()
        except Exception:
            _cleanup_running.clear()
            raise


# Metric name mapping for cleaner CSV output
METRIC_NAME_MAPPING = {
    'calls': 'calls',
    'exec_time_total': 'exec_time',
    'plan_time_total': 'plan_time',
    'rows': 'rows',
    'shared_bytes_hit_total': 'shared_blks_hit',
    'shared_bytes_read_total': 'shared_blks_read',
    'shared_bytes_dirtied_total': 'shared_blks_dirtied', 
    'shared_bytes_written_total': 'shared_blks_written',
    'block_read_total': 'blk_read_time',
    'block_write_total': 'blk_write_time'
}

def get_prometheus_client():
    """Get Prometheus client connection"""
    try:
        auth = None
        disable_ssl = True

        if os.environ.get('ENABLE_AMP', 'false').lower() == 'true':
            region = os.environ.get('AWS_REGION', 'us-east-1')
            service = 'aps'
            
            session = boto3.Session()
            credentials = session.get_credentials()
            
            if credentials:
                auth = AWS4Auth(
                    region=region,
                    service=service,
                    refreshable_credentials=credentials,
                )
            
            # Enable SSL verification for AMP
            disable_ssl = False

        vm_user = os.environ.get('VM_AUTH_USERNAME')
        vm_pass = os.environ.get('VM_AUTH_PASSWORD')
        if not auth and vm_user and vm_pass:
            auth = (vm_user, vm_pass)

        return PrometheusConnect(url=PROMETHEUS_URL, disable_ssl=disable_ssl, auth=auth)
    except Exception as e:
        logger.error(f"Failed to connect to Prometheus: {e}")
        raise


# Used to gate the cross-source-duplicate warning so we do not spam logs once
# per dashboard refresh. Keyed by datname; emits at most one warning per
# (datname, source-set) pair per process.
_logged_cross_source_warnings: set = set()


def _should_apply_db_filter(db_name: str | None) -> bool:
    """
    Decide whether ``db_name`` is a real database name worth filtering on.

    Skips empty strings, the literal Grafana ``All`` sentinel, and unresolved
    template variables (``$db_name``) that occasionally leak through when a
    dashboard panel is rendered before the variable is populated.
    """
    return bool(db_name) and db_name.lower() not in ('all', '') and not db_name.startswith('$')


def _build_query_text_lookup_sql(*, with_db_filter: bool, time_clause: str | None) -> str:
    """
    Build the SQL for queryid → query text lookups against the sink table.

    The sink table ``public.pgss_queryid_queries`` is LIST-partitioned by
    the ``dbname`` column, which stores the *pgwatch source name* (e.g.
    ``my-host-cone-demo``) — NOT the actual Postgres database name. The
    actual database name (``datname``) is stored inside the JSONB payload
    as ``data->>'real_dbname'``. Dashboards pass the Prometheus
    ``datname`` label, so the filter must be on the JSONB key. Filtering
    on the partition ``dbname`` column with a ``datname`` value yields
    zero rows — that was the root cause of MR !271.

    Args:
        with_db_filter: when True, append the
            ``data->>'real_dbname' = %s`` predicate (caller binds a
            ``datname``). The companion ``data ? 'real_dbname'`` is also
            emitted so the absent-key case is explicit (older pgwatch
            payloads or non-pgss rows missing the key are excluded
            deliberately, not by accident).
        time_clause: optional ``time > now() - ...`` fragment to bound
            the scan, e.g. ``"time > now() - make_interval(hours => %s)"``.

    When ``with_db_filter`` is True, the SQL also surfaces a
    ``source_dbnames`` column with the array of partition names that
    hold a row for the same queryid. Callers use this to detect (and
    log) the multi-source-ambiguity scenario flagged in MR !271's
    review: two pgwatch sources observing the same physical database
    would each contribute a row, and ``DISTINCT ON`` would pick one
    arbitrarily by ``time DESC``.

    The aggregation cannot use ``array_agg(distinct ...) OVER (...)``
    because PostgreSQL does not implement ``DISTINCT`` in window
    aggregates (`FeatureNotSupported`). Instead, the filtered scan
    is materialized once via a CTE and joined to a separate
    ``GROUP BY queryid`` aggregate — same single underlying scan
    plan, but only standard SQL constructs.

    Performance note: efficient evaluation of the
    ``data->>'real_dbname'`` predicate requires the expression index
    ``pgss_queryid_queries_real_dbname_time_idx`` declared in
    ``config/sink-postgres/init.sql`` (also covers the ``time``
    column). Without that index this becomes a per-partition seq scan.
    """
    where_clauses = [
        "data->>'queryid' is not null",
        "data->>'query' is not null",
    ]
    if with_db_filter:
        where_clauses.append("data ? 'real_dbname'")
        where_clauses.append("data->>'real_dbname' = %s")
    if time_clause:
        where_clauses.append(time_clause)
    where_sql = "\n              and ".join(where_clauses)

    if with_db_filter:
        return (
            "with filtered as (\n"
            "    select\n"
            "        data->>'queryid' as queryid,\n"
            "        data->>'query' as query,\n"
            "        dbname,\n"
            "        time\n"
            "    from public.pgss_queryid_queries\n"
            "    where\n"
            f"              {where_sql}\n"
            "),\n"
            "sources as (\n"
            "    select queryid, array_agg(distinct dbname) as source_dbnames\n"
            "    from filtered\n"
            "    group by queryid\n"
            "),\n"
            "latest as (\n"
            "    select distinct on (queryid) queryid, query\n"
            "    from filtered\n"
            "    order by queryid, time desc\n"
            ")\n"
            "select l.queryid, l.query, s.source_dbnames\n"
            "from latest l\n"
            "join sources s using (queryid)"
        )
    return (
        "select distinct on (data->>'queryid')\n"
        "    data->>'queryid' as queryid,\n"
        "    data->>'query' as query\n"
        "from public.pgss_queryid_queries\n"
        "where\n"
        f"              {where_sql}\n"
        "order by data->>'queryid', time desc"
    )


def _warn_on_cross_source_duplicates(db_name: str, duplicates: list) -> None:
    """
    Log a one-shot warning when a single ``datname`` is observed under
    more than one pgwatch source partition. ``DISTINCT ON`` silently
    picks one row per queryid by ``time DESC``, so without this signal a
    cross-tenant collision (two sources monitoring the same physical
    database name) would surface query texts from a non-deterministic
    source. Each duplicate is a tuple ``(queryid, source_dbnames)``.
    """
    if not duplicates:
        return
    # Dedupe by the full (datname, sorted source-set) tuple — the same
    # collision will be observed on every refresh, no need to log twice.
    fingerprint = (db_name, tuple(sorted({tuple(sorted(s[1])) for s in duplicates})))
    if fingerprint in _logged_cross_source_warnings:
        return
    _logged_cross_source_warnings.add(fingerprint)
    queryid_sample, sources_sample = duplicates[0]
    logger.warning(
        "Cross-source ambiguity for datname=%s: %d queryid(s) appear in "
        "multiple pgwatch source partitions; DISTINCT ON picks one by "
        "time DESC. Example: queryid=%s, sources=%s. Consider scoping "
        "by source name as well, or changing the dedup trigger to "
        "(real_dbname, queryid).",
        db_name, len(duplicates), queryid_sample, sorted(sources_sample),
    )


def get_query_texts_from_sink(db_name: str | None = None, truncation_mode: str = 'smart',
                              max_age_hours: int | None = None) -> dict:
    """
    Fetch queryid-to-query text mappings from the PostgreSQL sink database.

    Args:
        db_name: Optional Postgres database name (``datname``) to filter
            results. Matched against ``data->>'real_dbname'`` in the
            JSONB payload, NOT the partition ``dbname`` column (which
            holds the pgwatch source name).
        truncation_mode: 'smart' for smart truncation, 'raw' for simple truncation
        max_age_hours: Only return queryids seen within this many hours (None = use retention window)

    Returns:
        Dictionary mapping queryid to query text
    """
    if max_age_hours is None:
        max_age_hours = QUERYID_RETENTION_HOURS

    query_texts = {}

    conn = None
    try:
        conn = psycopg2.connect(POSTGRES_SINK_URL, connect_timeout=10)
        with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
            use_db_filter = _should_apply_db_filter(db_name)
            sql = _build_query_text_lookup_sql(
                with_db_filter=use_db_filter,
                time_clause='time > now() - make_interval(hours => %s)',
            )
            if use_db_filter:
                cursor.execute(sql, (db_name, max_age_hours))
            else:
                cursor.execute(sql, (max_age_hours,))

            cross_source_dupes = []
            for row in cursor:
                queryid = row['queryid']
                query_text = row['query']
                if use_db_filter:
                    sources = row['source_dbnames'] or []
                    if len(sources) > 1:
                        cross_source_dupes.append((queryid, list(sources)))
                if queryid:
                    if query_text:
                        if truncation_mode == 'raw':
                            normalized = ' '.join(query_text.split())
                            query_text = (normalized[:147] + '...') if len(normalized) > 150 else normalized
                        else:
                            query_text = smart_truncate_query(query_text, 150)
                    else:
                        query_text = ''
                    query_texts[queryid] = query_text

            if use_db_filter:
                _warn_on_cross_source_duplicates(db_name, cross_source_dupes)
    except Exception as e:
        logger.warning(f"Failed to fetch query texts from sink database: {e}")
    finally:
        if conn:
            conn.close()

    return query_texts


def read_version_file(filepath, default='unknown'):
    """Read version information from file"""
    try:
        with open(filepath, 'r') as f:
            return f.read().strip()
    except FileNotFoundError:
        return default


# Read version info at startup
APP_VERSION = read_version_file('/VERSION')
APP_BUILD_TS = read_version_file('/BUILD_TS')


@app.route('/version', methods=['GET'])
def version():
    """Return application version and build timestamp as array for Grafana Infinity datasource"""
    display = f"PostgresAI v{APP_VERSION} (built: {APP_BUILD_TS})"
    return jsonify([{
        "version": APP_VERSION,
        "build_ts": APP_BUILD_TS,
        "display": display
    }])


@app.route('/health', methods=['GET'])
def health():
    """Health check endpoint"""
    try:
        prom = get_prometheus_client()
        # Simple query to test connection
        prom.get_current_metric_value(metric_name='up')
        return jsonify({"status": "healthy", "prometheus_url": PROMETHEUS_URL})
    except Exception as e:
        return jsonify({"status": "unhealthy", "error": str(e)}), 500

@app.route('/pgss_metrics/csv', methods=['GET'])
def get_pgss_metrics_csv():
    """
    Get pg_stat_statements metrics as CSV with time-based difference calculation

    Query parameters:
    - time_start: Start time (ISO format or Unix timestamp)
    - time_end: End time (ISO format or Unix timestamp)
    - cluster_name: Cluster name filter (optional)
    - node_name: Node name filter (optional)
    - db_name: Database name filter (optional)
    - truncation_mode: 'smart' (default) or 'raw' for query text truncation
    """
    try:
        # Get query parameters
        time_start = request.args.get('time_start')
        time_end = request.args.get('time_end')
        cluster_name = request.args.get('cluster_name')
        node_name = request.args.get('node_name')
        truncation_mode = request.args.get('truncation_mode', 'smart')
        db_name = request.args.get('db_name')

        if not time_start or not time_end:
            return jsonify({"error": "time_start and time_end parameters are required"}), 400

        # Parse time parameters
        try:
            # Try parsing as Unix timestamp first
            start_dt = datetime.fromtimestamp(float(time_start), tz=timezone.utc)
        except ValueError:
            # Try parsing as ISO format
            start_dt = datetime.fromisoformat(time_start.replace('Z', '+00:00'))

        try:
            end_dt = datetime.fromtimestamp(float(time_end), tz=timezone.utc)
        except ValueError:
            end_dt = datetime.fromisoformat(time_end.replace('Z', '+00:00'))

        # Connect to Prometheus
        prom = get_prometheus_client()

        # Build the base query for pg_stat_statements metrics
        base_query = 'pgwatch_pg_stat_statements_calls'

        # Add filters if provided (escape values to prevent PromQL injection)
        filters = []
        if cluster_name:
            filters.append(f'cluster="{escape_promql_label(cluster_name)}"')
        if node_name:
            filters.append(f'instance=~".*{escape_promql_regex_literal(node_name)}.*"')
        if db_name:
            filters.append(f'datname="{escape_promql_label(db_name)}"')

        if filters:
            base_query += '{' + ','.join(filters) + '}'

        logger.info(f"Querying Prometheus with base query: {base_query}")

        # Get all pg_stat_statements metrics
        all_metrics = [
            'pgwatch_pg_stat_statements_calls',
            'pgwatch_pg_stat_statements_plans_total',
            'pgwatch_pg_stat_statements_exec_time_total',
            'pgwatch_pg_stat_statements_plan_time_total',
            'pgwatch_pg_stat_statements_rows',
            'pgwatch_pg_stat_statements_shared_bytes_hit_total',
            'pgwatch_pg_stat_statements_shared_bytes_read_total',
            'pgwatch_pg_stat_statements_shared_bytes_dirtied_total',
            'pgwatch_pg_stat_statements_shared_bytes_written_total',
            'pgwatch_pg_stat_statements_block_read_total',
            'pgwatch_pg_stat_statements_block_write_total',
            'pgwatch_pg_stat_statements_wal_records',
            'pgwatch_pg_stat_statements_wal_fpi',
            'pgwatch_pg_stat_statements_wal_bytes',
            'pgwatch_pg_stat_statements_temp_bytes_read',
            'pgwatch_pg_stat_statements_temp_bytes_written'
        ]

        # Apply filters to each metric
        filtered_metrics = []
        for metric in all_metrics:
            if filters:
                filtered_metrics.append(f'{metric}{{{",".join(filters)}}}')
            else:
                filtered_metrics.append(metric)

        # Get metrics at start and end times using instant queries
        start_data = []
        end_data = []

        for metric in filtered_metrics:
            try:
                start_metric_data = prom.get_metric_range_data(
                    metric_name=metric,
                    start_time=start_dt - timedelta(minutes=1),
                    end_time=start_dt + timedelta(minutes=1)
                )
                if start_metric_data:
                    start_data.extend(start_metric_data)

                end_metric_data = prom.get_metric_range_data(
                    metric_name=metric,
                    start_time=end_dt - timedelta(minutes=1),
                    end_time=end_dt + timedelta(minutes=1)
                )
                if end_metric_data:
                    end_data.extend(end_metric_data)
            except Exception as e:
                logger.warning(f"Failed to query metric {metric}: {e}")
                continue

        # Fetch query texts from sink database
        # Map legend_label values to truncation mode: displayname_raw_* -> raw, others -> smart
        trunc_mode = 'raw' if 'raw' in truncation_mode.lower() else 'smart'
        query_texts = get_query_texts_from_sink(db_name, truncation_mode=trunc_mode)
        logger.info(f"Fetched {len(query_texts)} query texts from sink database (mode: {trunc_mode})")

        # Process the data to calculate differences
        csv_data = process_pgss_data(start_data, end_data, start_dt, end_dt, query_texts)

        # Create CSV response
        output = io.StringIO()
        if csv_data:
            # Define explicit field order with queryid first, query_text second, then duration, then metrics with their rates
            base_fields = ['queryid', 'query_text', 'duration_seconds']
            all_metric_fields = []

            # Get metric fields from the mapping in specific order with their rates
            desired_order = [
                'calls', 'exec_time', 'plan_time', 'rows', 'shared_blks_hit',
                'shared_blks_read', 'shared_blks_dirtied', 'shared_blks_written',
                'blk_read_time', 'blk_write_time'
            ]

            for display_name in desired_order:
                if display_name in METRIC_NAME_MAPPING.values():
                    all_metric_fields.append(display_name)
                    all_metric_fields.append(f'{display_name}_per_sec')
                    all_metric_fields.append(f'{display_name}_per_call')

            # Combine all fields in desired order
            all_fields = base_fields + all_metric_fields
            
            writer = csv.DictWriter(output, fieldnames=all_fields)
            writer.writeheader()
            writer.writerows(csv_data)
        
        csv_content = output.getvalue()
        output.close()

        # Create response
        response = make_response(csv_content)
        response.headers['Content-Type'] = 'text/csv'
        response.headers['Content-Disposition'] = f'attachment; filename=pgss_metrics_{start_dt.strftime("%Y%m%d_%H%M%S")}_{end_dt.strftime("%Y%m%d_%H%M%S")}.csv'

        return response

    except Exception as e:
        logger.error(f"Error processing request: {e}")
        return jsonify({"error": str(e)}), 500

def process_pgss_data(start_data, end_data, start_time, end_time, query_texts=None):
    """
    Process pg_stat_statements data and calculate differences between start and end times

    Args:
        start_data: Prometheus data at start time
        end_data: Prometheus data at end time
        start_time: Start datetime
        end_time: End datetime
        query_texts: Optional dictionary mapping queryid to query text
    """
    if query_texts is None:
        query_texts = {}

    # Convert Prometheus data to dictionaries
    start_metrics = prometheus_to_dict(start_data, start_time)
    end_metrics = prometheus_to_dict(end_data, end_time)

    if not start_metrics and not end_metrics:
        return []

    # Create a combined dictionary with all unique query identifiers
    all_keys = set()
    all_keys.update(start_metrics.keys())
    all_keys.update(end_metrics.keys())

    result_rows = []

    # Calculate differences for each query
    for key in all_keys:
        start_metric = start_metrics.get(key, {})
        end_metric = end_metrics.get(key, {})

        # Extract identifier components from key
        db_name, query_id, user, instance = key

        # Calculate actual duration from metric timestamps
        start_timestamp = start_metric.get('timestamp')
        end_timestamp = end_metric.get('timestamp')

        if start_timestamp and end_timestamp:
            start_dt = datetime.fromisoformat(start_timestamp)
            end_dt = datetime.fromisoformat(end_timestamp)
            actual_duration = (end_dt - start_dt).total_seconds()
        else:
            # Fallback to query parameter duration if timestamps are missing
            actual_duration = (end_time - start_time).total_seconds()

        # Create result row with query text
        row = {
            'queryid': query_id,
            'query_text': query_texts.get(query_id, ''),
            'duration_seconds': actual_duration
        }

        # Numeric columns to calculate differences for (using original metric names)
        numeric_cols = list(METRIC_NAME_MAPPING.keys())
        
        # Calculate differences and rates
        for col in numeric_cols:
            start_val = start_metric.get(col, 0)
            end_val = end_metric.get(col, 0)
            diff = end_val - start_val
            
            # Use simplified display name for CSV columns
            display_name = METRIC_NAME_MAPPING[col]
            
            # Convert bytes to blocks for block-related metrics (PostgreSQL uses 8KB blocks)
            if 'blks' in display_name and 'bytes' in col:
                diff = diff / 8192  # Convert bytes to 8KB blocks
            
            row[display_name] = diff
            
            # Calculate rates per second
            if row['duration_seconds'] > 0:
                row[f'{display_name}_per_sec'] = diff / row['duration_seconds']
            else:
                row[f'{display_name}_per_sec'] = 0
                
            # Calculate per-call averages
            calls_diff = row.get('calls', 0)
            if calls_diff > 0:
                row[f'{display_name}_per_call'] = diff / calls_diff
            else:
                row[f'{display_name}_per_call'] = 0

        result_rows.append(row)

    # Sort by total execution time difference (descending)
    result_rows.sort(key=lambda x: x.get('exec_time', 0), reverse=True)

    return result_rows

def prometheus_to_dict(prom_data, timestamp):
    """
    Convert Prometheus API response to dictionary keyed by query identifiers
    """
    if not prom_data:
        return {}

    metrics_dict = {}

    for metric_data in prom_data:
        metric = metric_data.get('metric', {})
        values = metric_data.get('values', [])

        if not values:
            continue

        # Get the closest value to our timestamp
        closest_value = min(values, key=lambda x: abs(float(x[0]) - timestamp.timestamp()))

        # Create unique key for this query
        key = (
            metric.get('datname', ''),
            metric.get('queryid', ''),
            metric.get('user', ''),
            metric.get('instance', '')
        )

        # Initialize metric dict if not exists
        if key not in metrics_dict:
            metrics_dict[key] = {
                'timestamp': datetime.fromtimestamp(float(closest_value[0]), tz=timezone.utc).isoformat(),
            }

        # Add metric value
        metric_name = metric.get('__name__', 'pgwatch_pg_stat_statements_calls')
        clean_name = metric_name.replace('pgwatch_pg_stat_statements_', '')

        try:
            metrics_dict[key][clean_name] = float(closest_value[1])
        except (ValueError, IndexError):
            metrics_dict[key][clean_name] = 0

    return metrics_dict

@app.route('/metrics', methods=['GET'])
def list_metrics():
    """List available metrics in Prometheus"""
    try:
        prom = get_prometheus_client()
        metrics = prom.all_metrics()
        pgss_metrics = [m for m in metrics if 'pg_stat_statements' in m]
        return jsonify({"pg_stat_statements_metrics": pgss_metrics})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/debug/metrics', methods=['GET'])
def debug_metrics():
    """
    Debug endpoint to check what metrics are actually available in Prometheus
    """
    try:
        prom = get_prometheus_client()
        
        # Get all available metrics
        all_metrics = prom.all_metrics()
        
        # Filter for pg_btree_bloat metrics
        btree_metrics = [m for m in all_metrics if 'btree_bloat' in m]
        
        # Get sample data for each btree metric
        sample_data = {}
        for metric in btree_metrics[:5]:  # Limit to first 5 to avoid overwhelming
            try:
                result = prom.get_current_metric_value(metric_name=metric)
                sample_data[metric] = {
                    'count': len(result),
                    'sample_labels': [entry.get('metric', {}) for entry in result[:2]]  # First 2 entries
                }
            except Exception as e:
                sample_data[metric] = {'error': str(e)}
        
        return jsonify({
            'all_metrics_count': len(all_metrics),
            'btree_metrics': btree_metrics,
            'sample_data': sample_data
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/btree_bloat/csv', methods=['GET'])
def get_btree_bloat_csv():
    """
    Get the most recent pg_btree_bloat metrics as a CSV table.
    """
    try:
        # Get query parameters
        cluster_name = request.args.get('cluster_name')
        node_name = request.args.get('node_name')
        db_name = request.args.get('db_name')
        schemaname = request.args.get('schemaname')
        tblname = request.args.get('tblname')
        idxname = request.args.get('idxname')

        # Build label filters (escape values to prevent PromQL injection)
        _esc = escape_promql_label
        filters = []
        if cluster_name:
            filters.append(f'cluster="{_esc(cluster_name)}"')
        if node_name:
            filters.append(f'node_name="{_esc(node_name)}"')
        if schemaname:
            filters.append(f'schemaname="{_esc(schemaname)}"')
        if tblname:
            filters.append(f'tblname="{_esc(tblname)}"')
        if idxname:
            filters.append(f'idxname="{_esc(idxname)}"')
        if db_name:
            filters.append(f'datname="{_esc(db_name)}"')

        filter_str = '{' + ','.join(filters) + '}' if filters else ''

        # Metrics to fetch with last_over_time to get only the most recent value
        metric_queries = [
            f'last_over_time(pgwatch_pg_btree_bloat_real_size_mib{filter_str}[1d])',
            f'last_over_time(pgwatch_pg_btree_bloat_extra_size{filter_str}[1d])',
            f'last_over_time(pgwatch_pg_btree_bloat_extra_pct{filter_str}[1d])',
            f'last_over_time(pgwatch_pg_btree_bloat_fillfactor{filter_str}[1d])',
            f'last_over_time(pgwatch_pg_btree_bloat_bloat_size{filter_str}[1d])',
            f'last_over_time(pgwatch_pg_btree_bloat_bloat_pct{filter_str}[1d])',
            f'last_over_time(pgwatch_pg_btree_bloat_is_na{filter_str}[1d])',
        ]

        prom = get_prometheus_client()
        metric_results = {}

        for query in metric_queries:
            try:
                # Use custom_query instead of get_current_metric_value
                result = prom.custom_query(query=query)

                for entry in result:
                    metric_labels = entry.get('metric', {})
                    # Extract label values from Prometheus response (not user input —
                    # no PromQL escaping needed here; these go into the CSV output).
                    db_val = metric_labels.get('datname', '')
                    schema_val = metric_labels.get('schemaname', '')
                    tbl_val = metric_labels.get('tblname', '')
                    idx_val = metric_labels.get('idxname', '')
                    key = (db_val, schema_val, tbl_val, idx_val)

                    if key not in metric_results:
                        metric_results[key] = {
                            'database': db_val,
                            'schemaname': schema_val,
                            'tblname': tbl_val,
                            'idxname': idx_val,
                        }

                    # Extract metric type from query and store value
                    if 'real_size_mib' in query:
                        metric_results[key]['real_size_mib'] = float(entry['value'][1])
                    elif 'extra_size' in query and 'extra_pct' not in query:
                        metric_results[key]['extra_size'] = float(entry['value'][1])
                    elif 'extra_pct' in query:
                        metric_results[key]['extra_pct'] = float(entry['value'][1])
                    elif 'fillfactor' in query:
                        metric_results[key]['fillfactor'] = float(entry['value'][1])
                    elif 'bloat_size' in query:
                        metric_results[key]['bloat_size'] = float(entry['value'][1])
                    elif 'bloat_pct' in query:
                        metric_results[key]['bloat_pct'] = float(entry['value'][1])
                    elif 'is_na' in query:
                        metric_results[key]['is_na'] = int(float(entry['value'][1]))

            except Exception as e:
                logger.warning(f"Failed to query: {query}, error: {e}")
                continue

        # Prepare CSV output
        output = io.StringIO()
        fieldnames = [
            'database', 'schemaname', 'tblname', 'idxname',
            'real_size_mib', 'extra_size', 'extra_pct', 'fillfactor',
            'bloat_size', 'bloat_pct', 'is_na'
        ]
        writer = csv.DictWriter(output, fieldnames=fieldnames)
        writer.writeheader()
        for row in metric_results.values():
            writer.writerow(row)

        csv_content = output.getvalue()
        output.close()

        # Create response
        response = make_response(csv_content)
        response.headers['Content-Type'] = 'text/csv'
        response.headers['Content-Disposition'] = 'attachment; filename=btree_bloat_latest.csv'
        return response

    except Exception as e:
        logger.error(f"Error processing btree bloat request: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/table_info/csv', methods=['GET'])
def get_table_info_csv():
    """
    Get comprehensive table information including size metrics, tuple statistics, and I/O statistics as a CSV table.
    Supports both instant queries (without time parameters) and rate calculations over a time period.
    
    Query parameters:
    - time_start: Start time (ISO format or Unix timestamp) - optional
    - time_end: End time (ISO format or Unix timestamp) - optional
    - cluster_name: Cluster name filter (optional)
    - node_name: Node name filter (optional)
    - db_name: Database name filter (optional)
    - schemaname: Schema name filter (optional, supports regex with ~)
    - tblname: Table name filter (optional)
    """
    try:
        # Get query parameters
        time_start = request.args.get('time_start')
        time_end = request.args.get('time_end')
        cluster_name = request.args.get('cluster_name')
        node_name = request.args.get('node_name')
        db_name = request.args.get('db_name')
        schemaname = request.args.get('schemaname')
        tblname = request.args.get('tblname')

        # Determine if we should calculate rates
        calculate_rates = bool(time_start and time_end)
        
        if calculate_rates:
            # Parse time parameters
            try:
                start_dt = datetime.fromtimestamp(float(time_start), tz=timezone.utc)
            except ValueError:
                start_dt = datetime.fromisoformat(time_start.replace('Z', '+00:00'))

            try:
                end_dt = datetime.fromtimestamp(float(time_end), tz=timezone.utc)
            except ValueError:
                end_dt = datetime.fromisoformat(time_end.replace('Z', '+00:00'))

        # Build label filters (escape values to prevent PromQL injection)
        _esc = escape_promql_label
        filters = []
        if cluster_name:
            filters.append(f'cluster="{_esc(cluster_name)}"')
        if node_name:
            filters.append(f'node_name="{_esc(node_name)}"')
        if schemaname:
            # Use =~ with a regex-escaped literal so dots and other RE2 metacharacters
            # in schema names are matched literally, not as regex patterns.
            filters.append(f'schemaname=~"{escape_promql_regex_literal(schemaname)}"')
        if tblname:
            filters.append(f'tblname="{_esc(tblname)}"')
        if db_name:
            filters.append(f'datname="{_esc(db_name)}"')

        filter_str = '{' + ','.join(filters) + '}' if filters else ''

        prom = get_prometheus_client()
        
        # Define base metrics to query (without last_over_time wrapper for rate calculation)
        base_metrics = {
            # Size metrics
            'total_size': f'pgwatch_pg_class_total_relation_size_bytes{filter_str}',
            'table_size': f'pgwatch_table_size_detailed_table_main_size_b{filter_str}',
            'index_size': f'pgwatch_table_size_detailed_table_indexes_size_b{filter_str}',
            'toast_size': f'pgwatch_table_size_detailed_total_toast_size_b{filter_str}',
            # Scan statistics
            'seq_scan': f'pgwatch_pg_stat_all_tables_seq_scan{filter_str}',
            'idx_scan': f'pgwatch_pg_stat_all_tables_idx_scan{filter_str}',
            # Tuple statistics
            'n_tup_ins': f'pgwatch_table_stats_n_tup_ins{filter_str}',
            'n_tup_upd': f'pgwatch_table_stats_n_tup_upd{filter_str}',
            'n_tup_del': f'pgwatch_table_stats_n_tup_del{filter_str}',
            'n_tup_hot_upd': f'pgwatch_table_stats_n_tup_hot_upd{filter_str}',
            # I/O statistics
            'heap_blks_read': f'pgwatch_pg_statio_all_tables_heap_blks_read{filter_str}',
            'heap_blks_hit': f'pgwatch_pg_statio_all_tables_heap_blks_hit{filter_str}',
            'idx_blks_read': f'pgwatch_pg_statio_all_tables_idx_blks_read{filter_str}',
            'idx_blks_hit': f'pgwatch_pg_statio_all_tables_idx_blks_hit{filter_str}',
        }
        
        if calculate_rates:
            # Get metrics at start and end times
            start_data = {}
            end_data = {}
            
            for metric_name, metric_query in base_metrics.items():
                try:
                    # Get data at start time
                    start_result = prom.get_metric_range_data(
                        metric_name=metric_query,
                        start_time=start_dt - timedelta(minutes=1),
                        end_time=start_dt + timedelta(minutes=1)
                    )
                    if start_result:
                        start_data[metric_name] = start_result
                    
                    # Get data at end time
                    end_result = prom.get_metric_range_data(
                        metric_name=metric_query,
                        start_time=end_dt - timedelta(minutes=1),
                        end_time=end_dt + timedelta(minutes=1)
                    )
                    if end_result:
                        end_data[metric_name] = end_result
                except Exception as e:
                    logger.warning(f"Failed to query metric {metric_name}: {e}")
                    continue
            
            # Process the data to calculate rates
            metric_results = process_table_stats_with_rates(start_data, end_data, start_dt, end_dt)
        else:
            # Get instant values using last_over_time
            metric_results = {}
            for metric_name, metric_query in base_metrics.items():
                try:
                    result = prom.custom_query(query=f'last_over_time({metric_query}[1d])')
                    for entry in result:
                        metric_labels = entry.get('metric', {})
                        
                        # Use different key depending on label names
                        schema_label = metric_labels.get('schemaname') or metric_labels.get('schema', '')
                        table_label = metric_labels.get('relname') or metric_labels.get('table_name') or metric_labels.get('tblname', '')
                        
                        key = (
                            metric_labels.get('datname', ''),
                            schema_label,
                            table_label,
                        )
                        
                        if key not in metric_results:
                            metric_results[key] = {
                                'database': metric_labels.get('datname', ''),
                                'schema': schema_label,
                                'table_name': table_label,
                            }
                        
                        value = float(entry['value'][1])
                        metric_results[key][metric_name] = value
                except Exception as e:
                    logger.warning(f"Failed to query metric {metric_name}: {e}")
                    continue

        # Prepare CSV output
        output = io.StringIO()
        
        if calculate_rates:
            # Fields with rate calculations
            fieldnames = [
                'schema', 'table_name',
                # Size metrics (bytes)
                'total_size', 'table_size', 'index_size', 'toast_size',
                # Scan statistics with rates
                'seq_scans', 'seq_scans_per_sec',
                'idx_scans', 'idx_scans_per_sec',
                # Tuple statistics with rates
                'inserts', 'inserts_per_sec',
                'updates', 'updates_per_sec',
                'deletes', 'deletes_per_sec',
                'hot_updates', 'hot_updates_per_sec',
                # I/O statistics with rates (in bytes using block_size)
                'heap_blks_read', 'heap_blks_read_per_sec',
                'heap_blks_hit', 'heap_blks_hit_per_sec',
                'idx_blks_read', 'idx_blks_read_per_sec',
                'idx_blks_hit', 'idx_blks_hit_per_sec',
                'duration_seconds'
            ]
        else:
            # Fields without rate calculations
            fieldnames = [
                'schema', 'table_name',
                'total_size', 'table_size', 'index_size', 'toast_size',
                'seq_scan', 'idx_scan',
                'n_tup_ins', 'n_tup_upd', 'n_tup_del', 'n_tup_hot_upd',
                'heap_blks_read', 'heap_blks_hit',
                'idx_blks_read', 'idx_blks_hit'
            ]
            
            # Remove 'database' field from rows if present (not in fieldnames)
            for row in metric_results.values():
                row.pop('database', None)
        
        writer = csv.DictWriter(output, fieldnames=fieldnames)
        writer.writeheader()
        
        # Write rows (handle both dict and list)
        if isinstance(metric_results, dict):
            rows = metric_results.values()
        else:
            rows = metric_results
        
        for row in rows:
            writer.writerow(row)

        csv_content = output.getvalue()
        output.close()

        # Create response
        response = make_response(csv_content)
        response.headers['Content-Type'] = 'text/csv'
        
        if calculate_rates:
            filename = f'table_stats_{start_dt.strftime("%Y%m%d_%H%M%S")}_{end_dt.strftime("%Y%m%d_%H%M%S")}.csv'
        else:
            filename = 'table_stats_latest.csv'
        
        response.headers['Content-Disposition'] = f'attachment; filename={filename}'
        return response

    except Exception as e:
        logger.error(f"Error processing table stats request: {e}")
        return jsonify({"error": str(e)}), 500

def process_table_stats_with_rates(start_data, end_data, start_time, end_time):
    """
    Process table statistics and calculate rates between start and end times
    """
    # Convert data to dictionaries
    start_metrics = prometheus_table_to_dict(start_data, start_time)
    end_metrics = prometheus_table_to_dict(end_data, end_time)
    
    if not start_metrics and not end_metrics:
        return []
    
    # Get all unique table identifiers
    all_keys = set()
    all_keys.update(start_metrics.keys())
    all_keys.update(end_metrics.keys())
    
    result_rows = []
    
    for key in all_keys:
        start_metric = start_metrics.get(key, {})
        end_metric = end_metrics.get(key, {})
        
        # Extract identifier components from key
        db_name, schema_name, table_name = key
        
        # Calculate actual duration
        start_timestamp = start_metric.get('timestamp')
        end_timestamp = end_metric.get('timestamp')
        
        if start_timestamp and end_timestamp:
            start_dt = datetime.fromisoformat(start_timestamp)
            end_dt = datetime.fromisoformat(end_timestamp)
            actual_duration = (end_dt - start_dt).total_seconds()
        else:
            actual_duration = (end_time - start_time).total_seconds()
        
        # Create result row
        row = {
            'schema': schema_name,
            'table_name': table_name,
            'duration_seconds': actual_duration
        }
        
        # Counter metrics to calculate differences and rates
        counter_metrics = [
            'seq_scan', 'idx_scan', 'n_tup_ins', 'n_tup_upd', 
            'n_tup_del', 'n_tup_hot_upd', 'heap_blks_read', 'heap_blks_hit',
            'idx_blks_read', 'idx_blks_hit'
        ]
        
        # Mapping for display names
        display_names = {
            'seq_scan': 'seq_scans',
            'idx_scan': 'idx_scans',
            'n_tup_ins': 'inserts',
            'n_tup_upd': 'updates',
            'n_tup_del': 'deletes',
            'n_tup_hot_upd': 'hot_updates',
        }
        
        # Calculate differences and rates
        for metric in counter_metrics:
            start_val = start_metric.get(metric, 0)
            end_val = end_metric.get(metric, 0)
            diff = end_val - start_val
            
            # Use display name if available
            display_name = display_names.get(metric, metric)
            
            row[display_name] = diff
            
            # Calculate rate per second
            if actual_duration > 0:
                row[f'{display_name}_per_sec'] = diff / actual_duration
            else:
                row[f'{display_name}_per_sec'] = 0
        
        # Size metrics (just use end values, these don't need rates)
        for size_metric in ['total_size', 'table_size', 'index_size', 'toast_size']:
            row[size_metric] = end_metric.get(size_metric, 0)
        
        result_rows.append(row)
    
    # Sort by total size descending
    result_rows.sort(key=lambda x: x.get('total_size', 0), reverse=True)
    
    return result_rows

def prometheus_table_to_dict(prom_data, timestamp):
    """
    Convert Prometheus table metrics to dictionary keyed by table identifiers
    """
    if not prom_data:
        return {}
    
    metrics_dict = {}
    
    for metric_name, metric_results in prom_data.items():
        for metric_data in metric_results:
            metric = metric_data.get('metric', {})
            values = metric_data.get('values', [])
            
            if not values:
                continue
            
            # Get the closest value to our timestamp
            closest_value = min(values, key=lambda x: abs(float(x[0]) - timestamp.timestamp()))
            
            # Handle different label names
            schema_label = metric.get('schemaname') or metric.get('schema', '')
            table_label = metric.get('relname') or metric.get('table_name') or metric.get('tblname', '')
            
            # Create unique key for this table
            key = (
                metric.get('datname', ''),
                schema_label,
                table_label,
            )
            
            # Initialize metric dict if not exists
            if key not in metrics_dict:
                metrics_dict[key] = {
                    'timestamp': datetime.fromtimestamp(float(closest_value[0]), tz=timezone.utc).isoformat(),
                }
            
            # Add metric value
            try:
                metrics_dict[key][metric_name] = float(closest_value[1])
            except (ValueError, IndexError):
                metrics_dict[key][metric_name] = 0
    
    return metrics_dict


@app.route('/query_texts', methods=['GET'])
def get_query_texts():
    """
    Get queryid-to-query text mappings for use in Grafana TopN chart legends.

    Returns a JSON array of objects with 'queryid', 'query_text', and 'displayName' fields,
    suitable for use with Grafana's Infinity datasource and "Config from query results" transformation.

    The 'displayName' field contains a smart-truncated query string suitable for chart legends:
    - Leading comments (/* ... */ and -- ...) are stripped
    - SELECT queries show: "SELECT ... FROM <table_names>"
    - CTEs show: "WITH cte1, cte2 SELECT ... FROM ..."
    - INSERT/UPDATE/DELETE show the target table name
    - Falls back to simple truncation on parse errors

    Query parameters:
    - db_name: Database name filter (optional)
    - truncate: Max characters for query text (default: 40 for chart legends)
    """
    try:
        db_name = request.args.get('db_name')

        # Validate truncate parameter
        try:
            truncate_len = int(request.args.get('truncate', 40))
            if truncate_len < 1 or truncate_len > 10000:
                return jsonify({'error': 'truncate must be between 1 and 10000'}), 400
        except (ValueError, TypeError):
            return jsonify({'error': 'truncate must be a valid integer'}), 400

        query_texts = {}

        conn = None
        try:
            conn = psycopg2.connect(POSTGRES_SINK_URL, connect_timeout=10)
            with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
                use_db_filter = _should_apply_db_filter(db_name)
                # Bound the scan by ``QUERYID_RETENTION_HOURS`` — without
                # a time predicate every dashboard load scans every
                # partition end-to-end, which degrades as data accumulates.
                sql = _build_query_text_lookup_sql(
                    with_db_filter=use_db_filter,
                    time_clause='time > now() - make_interval(hours => %s)',
                )
                if use_db_filter:
                    cursor.execute(sql, (db_name, QUERYID_RETENTION_HOURS))
                else:
                    cursor.execute(sql, (QUERYID_RETENTION_HOURS,))

                cross_source_dupes = []
                for row in cursor:
                    queryid = row['queryid']
                    query_text = row['query']
                    if use_db_filter:
                        sources = row['source_dbnames'] or []
                        if len(sources) > 1:
                            cross_source_dupes.append((queryid, list(sources)))
                    if queryid:
                        # Use defensive check for None (despite SQL filter, drivers may return None)
                        query_text = smart_truncate_query(query_text or '', truncate_len)
                        query_texts[queryid] = query_text or ''

                if use_db_filter:
                    _warn_on_cross_source_duplicates(db_name, cross_source_dupes)
        except Exception as e:
            logger.warning(f"Failed to fetch query texts from sink database: {e}")
        finally:
            if conn:
                conn.close()

        # Return as JSON array for Grafana Infinity datasource
        # Include 'displayName' for use with "Config from query results" transformation
        # displayName = query_text for legend series name (readable)
        # queryid remains separate for navigation links to Dashboard 3
        result = [
            {
                'queryid': qid,
                'query_text': qtext,
                'displayName': qtext if qtext else qid
            }
            for qid, qtext in query_texts.items()
        ]

        return jsonify(result)

    except Exception as e:
        logger.error(f"Error fetching query texts: {e}")
        return jsonify({"error": str(e)}), 500


def _escape_prometheus_label(value):
    """Escape a string for use as a Prometheus label value.

    Handles:
    - Backslashes: \\ -> \\\\
    - Double quotes: " -> \\"
    - Newlines, carriage returns, tabs: replaced with space
    """
    if not value:
        return ""
    return (value
            .replace('\\', '\\\\')
            .replace('"', '\\"')
            .replace('\n', ' ')
            .replace('\r', ' ')
            .replace('\t', ' '))


@app.route('/query_info_metrics', methods=['GET'])
def get_query_info_metrics():
    """
    Return query info metrics in Prometheus exposition format.

    This endpoint is designed to be scraped by VictoriaMetrics/Prometheus
    to create pgwatch_query_info{queryid="...", displayname="...", ...} metrics
    that can be used with group_left() to add display names to chart legends.

    The metrics include multiple truncation levels for dashboard flexibility:
    - displayname: Short version (30 chars) - default for legends
    - displayname_medium: Medium version (60 chars)
    - displayname_long: Long version (100 chars)

    Query parameters:
    - db_name: Database name filter (optional)
    """
    try:
        db_name = request.args.get('db_name')

        # Truncation lengths for different display modes
        TRUNCATE_SHORT = 30
        TRUNCATE_MEDIUM = 60
        TRUNCATE_LONG = 100

        query_data = {}

        conn = None
        try:
            conn = psycopg2.connect(POSTGRES_SINK_URL, connect_timeout=10)
            # Run retention cleanup in a background thread so it doesn't
            # add latency to every Prometheus scrape.
            try:
                _start_retention_cleanup_thread()
            except Exception as e:
                logger.warning("Failed to start retention cleanup thread: %s", e)

            with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
                # Only export queryids recently seen in pg_stat_statements.
                # The dedup trigger refreshes the timestamp on each collection,
                # so active queryids have time within the last few minutes.
                use_db_filter = _should_apply_db_filter(db_name)
                sql = _build_query_text_lookup_sql(
                    with_db_filter=use_db_filter,
                    time_clause='time > now() - make_interval(mins => %s)',
                )
                if use_db_filter:
                    cursor.execute(sql, (db_name, QUERYID_ACTIVE_MINUTES))
                else:
                    cursor.execute(sql, (QUERYID_ACTIVE_MINUTES,))

                cross_source_dupes = []
                for row in cursor:
                    queryid = row['queryid']
                    query_text = row['query'] or ''  # Defensive check for None
                    if use_db_filter:
                        sources = row['source_dbnames'] or []
                        if len(sources) > 1:
                            cross_source_dupes.append((queryid, list(sources)))
                    if queryid:
                        # Normalize whitespace for raw truncation
                        normalized_text = ' '.join(query_text.split()) if query_text else ''

                        # Smart truncation (extracts table names, CTEs, etc.)
                        short_text = smart_truncate_query(query_text, TRUNCATE_SHORT)
                        medium_text = smart_truncate_query(query_text, TRUNCATE_MEDIUM)
                        long_text = smart_truncate_query(query_text, TRUNCATE_LONG)

                        # Format: "queryid | query text" - queryid prefix for easy copy-paste
                        short = _escape_prometheus_label(f"{queryid} | {short_text}" if short_text else queryid)
                        medium = _escape_prometheus_label(f"{queryid} | {medium_text}" if medium_text else queryid)
                        long = _escape_prometheus_label(f"{queryid} | {long_text}" if long_text else queryid)

                        # Raw truncation (simple character limit, no smart extraction)
                        raw_short_text = (normalized_text[:TRUNCATE_SHORT-3] + '...') if len(normalized_text) > TRUNCATE_SHORT else normalized_text
                        raw_medium_text = (normalized_text[:TRUNCATE_MEDIUM-3] + '...') if len(normalized_text) > TRUNCATE_MEDIUM else normalized_text
                        raw_long_text = (normalized_text[:TRUNCATE_LONG-3] + '...') if len(normalized_text) > TRUNCATE_LONG else normalized_text

                        raw_short = _escape_prometheus_label(f"{queryid} | {raw_short_text}" if raw_short_text else queryid)
                        raw_medium = _escape_prometheus_label(f"{queryid} | {raw_medium_text}" if raw_medium_text else queryid)
                        raw_long = _escape_prometheus_label(f"{queryid} | {raw_long_text}" if raw_long_text else queryid)

                        # Full query text (no truncation) - limit to 500 chars for Prometheus label safety
                        full_text = normalized_text[:500] if normalized_text else ''
                        full = _escape_prometheus_label(f"{queryid} | {full_text}" if full_text else queryid)

                        # Just queryid (for users who prefer numeric IDs)
                        queryid_only = _escape_prometheus_label(queryid)

                        query_data[queryid] = {
                            'short': short,
                            'medium': medium,
                            'long': long,
                            'raw_short': raw_short,
                            'raw_medium': raw_medium,
                            'raw_long': raw_long,
                            'full': full,
                            'queryid_only': queryid_only,
                        }

                if use_db_filter:
                    _warn_on_cross_source_duplicates(db_name, cross_source_dupes)
                logger.info(f"Exported {len(query_data)} active queryids for metrics")
        except Exception as e:
            logger.warning(f"Failed to fetch query texts from sink database: {e}")
        finally:
            if conn:
                conn.close()

        # Generate Prometheus exposition format
        lines = ['# HELP pgwatch_query_info Query ID to display name mapping for chart legends']
        lines.append('# TYPE pgwatch_query_info gauge')

        for queryid, names in query_data.items():
            # Include all truncation levels as labels
            lines.append(
                f'pgwatch_query_info{{queryid="{queryid}",'
                f'displayname="{names["short"]}",'
                f'displayname_medium="{names["medium"]}",'
                f'displayname_long="{names["long"]}",'
                f'displayname_raw_short="{names["raw_short"]}",'
                f'displayname_raw_medium="{names["raw_medium"]}",'
                f'displayname_raw_long="{names["raw_long"]}",'
                f'displayname_full="{names["full"]}",'
                f'displayname_queryid="{names["queryid_only"]}"}} 1'
            )

        response = make_response('\n'.join(lines) + '\n')
        response.headers['Content-Type'] = 'text/plain; charset=utf-8'
        return response

    except Exception as e:
        logger.error(f"Error generating query info metrics: {e}")
        return f"# Error: {str(e)}\n", 500


def _check_debug_auth():
    """Validate debug endpoint auth. Returns None on success, a (response, status) tuple on failure."""
    if not _DEBUG_SECRET_KEY:
        return jsonify({'error': 'Debug auth not configured (DEBUG_SECRET_KEY env var not set)'}), 403
    auth_header = request.headers.get('Authorization', '')
    if not hmac.compare_digest(auth_header, f'Bearer {_DEBUG_SECRET_KEY}'):
        return jsonify({'error': 'Unauthorized'}), 401
    return None


@app.route('/execute-query', methods=['POST'])
def execute_query():
    """Debug-only endpoint for executing read-only queries against the configured database.

    Requires ENABLE_DEBUG=true and a valid Authorization: Bearer <DEBUG_SECRET_KEY> header.
    Returns 404 when debug mode is disabled so the endpoint is not discoverable in production.
    """
    if os.environ.get('ENABLE_DEBUG', 'false').lower() != 'true':
        return jsonify({'error': 'Not found'}), 404

    auth_error = _check_debug_auth()
    if auth_error:
        return auth_error

    data = request.get_json(silent=True)
    if not data:
        return jsonify({'error': 'Request body must be JSON with Content-Type: application/json'}), 400

    query = data.get('query')
    if not query or not isinstance(query, str):
        return jsonify({'error': 'query field is required and must be a string'}), 400

    # Strip leading block comments before allowlist check so that
    # "/* comment */ SELECT ..." is still accepted.
    query_stripped = re.sub(r'^(\s*/\*.*?\*/\s*)+', '', query, flags=re.DOTALL).strip()
    if not _SQL_ALLOWLIST_RE.match(query_stripped):
        return jsonify({'error': 'Only SELECT, EXPLAIN, and SHOW statements are permitted'}), 400

    # Strip ALL comments (block and line) before blocklist / multi-statement
    # checks.  The helper is quote-aware: comment delimiters that appear
    # inside single-quoted string literals are left intact, preventing bypass
    # attacks like:  SELECT 'x' || '--' || (SELECT * FROM pg_shadow)
    query_no_comments = _strip_sql_comments(query_stripped)
    if _SQL_BLOCKLIST_RE.search(query_no_comments):
        return jsonify({'error': 'Query contains disallowed functions or clauses '
                        '(pg_read_file, pg_read_binary_file, pg_ls_dir, pg_stat_file, '
                        'pg_shadow, pg_authid, WITH)'}), 400

    # Reject multi-statement queries (reuses query_no_comments already computed above).
    if ';' in query_no_comments:
        return jsonify({'error': 'Multi-statement queries are not permitted'}), 400

    conn = None
    try:
        cfg = {k: v for k, v in DB_CONFIG.items() if v is not None}
        conn = psycopg2.connect(**cfg, connect_timeout=10)
        conn.set_session(autocommit=True, readonly=True)
        with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
            cursor.execute('SET statement_timeout = %s', ('10s',))
            cursor.execute(query)
            rows = cursor.fetchall()
            columns = [desc[0] for desc in cursor.description] if cursor.description else []
        return jsonify({'columns': columns, 'rows': [list(row) for row in rows]})
    except Exception as e:
        logger.error(f"Debug execute-query error: {e}")
        return jsonify({'error': 'Query execution failed'}), 500
    finally:
        if conn:
            conn.close()


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False)
