import psycopg2
from psycopg2.pool import ThreadedConnectionPool
import json
import redis
import logging
import os
import threading
from contextvars import ContextVar
from time import perf_counter
from quillsql.error import PgQueryError
from quillsql.db.db_helper import connect_to_db, disconnect_from_db, run_query_by_db

# The TTL for new cache entries (default: 1h)
DEFAULT_CACHE_TTL = 24 * 60 * 60
DEFAULT_POSTGRES_POOL_MIN = 4
DEFAULT_POSTGRES_POOL_MAX = 24
DEFAULT_RUN_QUERY_MAX_WORKERS = 8
DEFAULT_DB_MAX_INFLIGHT = 12
DEFAULT_ACQUIRE_WAIT_EWMA_ALPHA = 0.2
_REQUEST_ID = ContextVar("quillsql_request_id", default=None)


class CachedConnection:
    _global_db_lock = threading.Lock()
    _global_db_semaphore = None
    _global_db_cap = None
    _acquire_wait_lock = threading.Lock()
    _acquire_wait_ewma_ms = None
    _pool_acquire_wait_lock = threading.Lock()
    _pool_acquire_wait_ewma_ms = None

    def __init__(self, database_type, config, cache_config, using_connection_string):
        self.database_type = database_type
        self.using_connection_string = using_connection_string
        self.config = config
        self.profile_logging = self._is_profile_enabled(cache_config)
        self.profile_logger = logging.getLogger("quillsql.profile")
        self.db_max_inflight = self._resolve_db_max_inflight(cache_config)
        self.acquire_wait_alpha = self._resolve_acquire_wait_alpha(cache_config)
        self._ensure_global_db_semaphore(self.db_max_inflight)
        self.connection = connect_to_db(database_type, config, using_connection_string)
        self.postgres_pool = self._create_postgres_pool(cache_config)
        self._pool_lock = threading.Lock()
        self._pool_in_use = 0
        self.cache = self.get_cache(cache_config)
        self.ttl = cache_config and cache_config.get("ttl") or DEFAULT_CACHE_TTL
        self.orgId = None

    def _resolve_db_max_inflight(self, cache_config):
        configured = cache_config and cache_config.get("db_max_inflight")
        if configured is None:
            configured = os.getenv("QUILL_DB_MAX_INFLIGHT", DEFAULT_DB_MAX_INFLIGHT)
        return max(1, int(configured))

    def _resolve_acquire_wait_alpha(self, cache_config):
        configured = cache_config and cache_config.get("acquire_wait_ewma_alpha")
        if configured is None:
            configured = os.getenv(
                "QUILL_ACQUIRE_WAIT_EWMA_ALPHA", DEFAULT_ACQUIRE_WAIT_EWMA_ALPHA
            )
        alpha = float(configured)
        return min(1.0, max(0.01, alpha))

    @classmethod
    def _ensure_global_db_semaphore(cls, cap):
        with cls._global_db_lock:
            if cls._global_db_semaphore is None or cls._global_db_cap != cap:
                cls._global_db_semaphore = threading.BoundedSemaphore(cap)
                cls._global_db_cap = cap

    def _acquire_global_db_slot(self):
        started = perf_counter()
        self._global_db_semaphore.acquire()
        wait_ms = round((perf_counter() - started) * 1000, 2)
        cls = type(self)
        with cls._acquire_wait_lock:
            prev = cls._acquire_wait_ewma_ms
            if prev is None:
                cls._acquire_wait_ewma_ms = wait_ms
            else:
                cls._acquire_wait_ewma_ms = (
                    (1 - self.acquire_wait_alpha) * prev
                    + self.acquire_wait_alpha * wait_ms
                )
        self._profile_log(
            "global_db_slot_acquired",
            wait_ms=wait_ms,
            db_max_inflight=self._global_db_cap,
            acquire_wait_ewma_ms=(
                round(self.get_recent_acquire_wait_ms(), 2)
                if self.get_recent_acquire_wait_ms() is not None
                else None
            ),
        )

    def _release_global_db_slot(self):
        self._global_db_semaphore.release()

    def _is_profile_enabled(self, cache_config):
        if cache_config and cache_config.get("profile_logging") is not None:
            return bool(cache_config.get("profile_logging"))
        enabled = str(os.getenv("QUILL_PROFILE_LOGGING", "")).strip().lower()
        return enabled in ("1", "true", "yes", "on")

    def _profile_log(self, event, **fields):
        if not self.profile_logging:
            return
        request_id = _REQUEST_ID.get()
        payload = {"component": "cached_connection", "event": event, **fields}
        if request_id:
            payload["request_id"] = request_id
        self.profile_logger.info("[quillsql.profile] %s", json.dumps(payload, default=str))

    def _is_postgres(self):
        return self.database_type and self.database_type.lower() in ("postgresql", "postgres")

    def _create_postgres_pool(self, cache_config):
        if not self._is_postgres():
            return None

        pool_min = cache_config and cache_config.get("postgres_pool_min")
        pool_max = cache_config and cache_config.get("postgres_pool_max")
        if pool_min is None:
            pool_min = os.getenv("QUILL_POSTGRES_POOL_MIN", DEFAULT_POSTGRES_POOL_MIN)
        if pool_max is None:
            pool_max = os.getenv("QUILL_POSTGRES_POOL_MAX", DEFAULT_POSTGRES_POOL_MAX)

        pool_min = max(1, int(pool_min))
        pool_max = max(pool_min, int(pool_max))

        self._profile_log(
            "postgres_pool_config",
            pool_min=pool_min,
            pool_max=pool_max,
        )

        if self.using_connection_string:
            return ThreadedConnectionPool(pool_min, pool_max, self.config)

        return ThreadedConnectionPool(
            pool_min,
            pool_max,
            database=self.config["dbname"],
            user=self.config["user"],
            password=self.config["password"],
            host=self.config["host"],
            port=self.config["port"],
        )

    def get_cache(self, cache_config):
        cache_type = cache_config and cache_config.get("cache_type")
        if cache_type and cache_type == "redis" or cache_type == "rediss":
            return redis.Redis(
                host=cache_config.get("host", "localhost"),
                port=cache_config.get("port", 6379),
                username=cache_config.get("username", "default"),
                password=cache_config.get("password"),
            )
        return None

    def _is_connection_closed(self):
        return self.connection is None or getattr(self.connection, "closed", True)

    def _open_connection(self):
        self.connection = connect_to_db(
            self.database_type, self.config, self.using_connection_string
        )
        return self.connection

    def _close_connection(self):
        if self.connection is None:
            return
        try:
            self.connection.close()
        except Exception:
            pass
        finally:
            self.connection = None

    def exec_with_reconnect(self, sql):
        reconnect_count = 0
        last_error = None
        while reconnect_count < 10:
            try:
                if self._is_connection_closed():
                    self._open_connection()
                return run_query_by_db(self.database_type, sql, self.connection)
            except psycopg2.Error as err:
                last_error = err
                reconnect_count += 1
                self._close_connection()
            except Exception:
                self._close_connection()
                raise
        diag = getattr(last_error, "diag", None)
        position = getattr(diag, "statement_position", None)
        raise PgQueryError(last_error or "Failed to execute query", sql, position)

    def exec(self, sql):
        # Use isolated connections for Postgres queries so concurrent SDK calls
        # on the same Quill instance do not contend on a shared connection.
        if self._is_postgres():
            return self.exec_with_new_connection(sql)
        try:
            started = perf_counter()
            return self.exec_with_reconnect(sql)
        except psycopg2.Error as err:
            raise PgQueryError(err, sql, err.diag.statement_position)
        finally:
            self._profile_log(
                "db_exec_shared_connection_done",
                database_type=self.database_type,
                elapsed_ms=round((perf_counter() - started) * 1000, 2),
                sql_preview=str(sql).strip().replace("\n", " ")[:180],
            )


    def exec_with_new_connection(self, sql):
        self._acquire_global_db_slot()
        try:
            if self.postgres_pool:
                connection = None
                started = perf_counter()
                query_started = None
                try:
                    connection = self.postgres_pool.getconn()
                    with self._pool_lock:
                        self._pool_in_use += 1
                    self._profile_log(
                        "pool_connection_acquired",
                        database_type=self.database_type,
                        elapsed_ms=round((perf_counter() - started) * 1000, 2),
                    )
                    self._update_pool_acquire_wait_ewma(
                        round((perf_counter() - started) * 1000, 2)
                    )
                    query_started = perf_counter()
                    return run_query_by_db(self.database_type, sql, connection)
                except psycopg2.Error as err:
                    raise PgQueryError(err, sql, err.diag.statement_position)
                finally:
                    if connection is not None:
                        should_close = bool(getattr(connection, "closed", False))
                        self.postgres_pool.putconn(connection, close=should_close)
                        with self._pool_lock:
                            self._pool_in_use = max(0, self._pool_in_use - 1)
                    self._profile_log(
                        "db_exec_pooled_connection_done",
                        database_type=self.database_type,
                        elapsed_ms=round((perf_counter() - started) * 1000, 2),
                        sql_exec_ms=(
                            round((perf_counter() - query_started) * 1000, 2)
                            if query_started is not None
                            else None
                        ),
                        sql_preview=str(sql).strip().replace("\n", " ")[:180],
                    )

            connection = connect_to_db(
                self.database_type, self.config, self.using_connection_string
            )
            started = perf_counter()
            try:
                return run_query_by_db(self.database_type, sql, connection)
            except psycopg2.Error as err:
                raise PgQueryError(err, sql, err.diag.statement_position)
            finally:
                disconnect_from_db(self.database_type, connection)
                self._profile_log(
                    "db_exec_new_connection_done",
                    database_type=self.database_type,
                    elapsed_ms=round((perf_counter() - started) * 1000, 2),
                    sql_preview=str(sql).strip().replace("\n", " ")[:180],
                )
        finally:
            self._release_global_db_slot()

    def __del__(self):
        if getattr(self, "postgres_pool", None):
            try:
                self.postgres_pool.closeall()
            except Exception:
                pass

    def query_with_new_connection(self, sql):
        if not self.cache:
            return self.exec_with_new_connection(sql)

        key = f"{self.orgId}:{sql}"
        started = perf_counter()
        cached_result = self.cache.get(key)
        if cached_result:
            cached = json.loads(cached_result)
            self._profile_log(
                "cache_hit",
                cache_key_prefix=str(key)[:120],
                elapsed_ms=round((perf_counter() - started) * 1000, 2),
            )
            return cached
        else:
            new_result = self.exec_with_new_connection(sql)
            new_result_string = json.dumps(new_result)
            self.cache.set(key, new_result_string, "EX", DEFAULT_CACHE_TTL)
            self._profile_log(
                "cache_miss_set",
                cache_key_prefix=str(key)[:120],
                elapsed_ms=round((perf_counter() - started) * 1000, 2),
            )
            return new_result

    def query(self, sql):
        if not self.cache:
            return self.exec(sql)

        key = f"{self.orgId}:{sql}"
        started = perf_counter()
        cached_result = self.cache.get(key)
        if cached_result:
            cached = json.loads(cached_result)
            self._profile_log(
                "cache_hit",
                cache_key_prefix=str(key)[:120],
                elapsed_ms=round((perf_counter() - started) * 1000, 2),
            )
            return cached
        else:
            new_result = self.exec(sql)
            new_result_string = json.dumps(new_result)
            self.cache.set(key, new_result_string, "EX", DEFAULT_CACHE_TTL)
            self._profile_log(
                "cache_miss_set",
                cache_key_prefix=str(key)[:120],
                elapsed_ms=round((perf_counter() - started) * 1000, 2),
            )
            return new_result

    def set_request_id(self, request_id):
        return _REQUEST_ID.set(request_id)

    def reset_request_id(self, token):
        _REQUEST_ID.reset(token)

    def get_pool_max_connections(self):
        if not self.postgres_pool:
            return None
        return int(getattr(self.postgres_pool, "maxconn", 0) or 0)

    def get_pool_pressure(self):
        pool_max = self.get_pool_max_connections()
        if not pool_max:
            return None
        with self._pool_lock:
            in_use = self._pool_in_use
        ratio = in_use / float(pool_max)
        return {
            "in_use": in_use,
            "max": pool_max,
            "ratio": ratio,
        }

    def get_recent_acquire_wait_ms(self):
        cls = type(self)
        with cls._acquire_wait_lock:
            return cls._acquire_wait_ewma_ms

    def _update_pool_acquire_wait_ewma(self, wait_ms):
        cls = type(self)
        with cls._pool_acquire_wait_lock:
            prev = cls._pool_acquire_wait_ewma_ms
            if prev is None:
                cls._pool_acquire_wait_ewma_ms = wait_ms
            else:
                cls._pool_acquire_wait_ewma_ms = (
                    (1 - self.acquire_wait_alpha) * prev
                    + self.acquire_wait_alpha * wait_ms
                )

    def get_recent_pool_acquire_wait_ms(self):
        cls = type(self)
        with cls._pool_acquire_wait_lock:
            return cls._pool_acquire_wait_ewma_ms

    def get_run_query_max_workers(self, run_query_config=None):
        config_value = run_query_config and run_query_config.get("maxWorkers")
        env_value = os.getenv("QUILL_RUN_QUERY_MAX_WORKERS")
        candidate = config_value if config_value is not None else env_value
        if candidate is None:
            candidate = DEFAULT_RUN_QUERY_MAX_WORKERS
        workers = max(1, int(candidate))
        pool_max = self.get_pool_max_connections()
        if pool_max:
            workers = min(workers, pool_max)
        return workers
