"""Worker-local lazy SQLAlchemy Engine registry."""

import threading
import time
from contextlib import contextmanager
from dataclasses import dataclass, field

from sqlalchemy import event
from sqlalchemy.engine import Engine

from app.core.data_source.circuit_breaker import CircuitBreaker
from app.core.data_source.models import PoolKey, PoolStatus


@dataclass
class PoolEntry:
    key: PoolKey
    engine: object
    breaker: CircuitBreaker
    created_at: float
    last_used_at: float
    leases: int = 0
    draining: bool = False
    drain_deadline: float = 0.0
    checkout_wait_ms: float = 0.0
    metrics: dict = field(
        default_factory=lambda: {
            "connection_created_total": 0,
            "connection_failed_total": 0,
            "pool_timeout_total": 0,
            "invalidated_total": 0,
            "checkout_total": 0,
            "checkin_total": 0,
            "checked_out_peak": 0,
            "query_total": 0,
            "last_query_duration_ms": 0.0,
        }
    )


class PoolRegistry:
    def __init__(self, settings, *, clock=None):
        self.settings = dict(settings)
        self._clock = clock or time.monotonic
        self._lock = threading.RLock()
        self._entries = {}
        self._draining = []

    def _attach_metrics(self, entry):
        if not isinstance(entry.engine, Engine):
            return

        def increment(name):
            with self._lock:
                entry.metrics[name] += 1

        def record_checkout(*_args):
            with self._lock:
                entry.metrics["checkout_total"] += 1
                entry.metrics["checked_out_peak"] = max(
                    entry.metrics["checked_out_peak"],
                    self._pool_value(entry.engine, "checkedout"),
                )

        event.listen(
            entry.engine,
            "connect",
            lambda *_args: increment("connection_created_total"),
        )
        event.listen(
            entry.engine,
            "checkout",
            record_checkout,
        )
        event.listen(
            entry.engine,
            "checkin",
            lambda *_args: increment("checkin_total"),
        )
        event.listen(
            entry.engine,
            "invalidate",
            lambda *_args: increment("invalidated_total"),
        )

    def _dispose(self, entry):
        if getattr(entry.engine, "_dataops_disposed", False):
            return
        entry.engine.dispose()
        try:
            setattr(entry.engine, "_dataops_disposed", True)
        except Exception:
            pass

    def _start_draining_locked(self, entry):
        if entry.draining:
            return
        self._entries.pop(entry.key, None)
        entry.draining = True
        entry.drain_deadline = self._clock() + float(
            self.settings.get("drain_timeout", 30)
        )
        entry.metrics["invalidated_total"] += 1
        if entry.leases == 0:
            self._dispose(entry)
        else:
            self._draining.append(entry)

    def _entry_locked(self, key, engine_factory):
        entry = self._entries.get(key)
        if entry is not None:
            return entry
        for old in list(self._entries.values()):
            if old.key.data_source_uid == key.data_source_uid:
                self._start_draining_locked(old)
        now = self._clock()
        entry = PoolEntry(
            key=key,
            engine=engine_factory(),
            breaker=CircuitBreaker(clock=self._clock),
            created_at=now,
            last_used_at=now,
        )
        self._entries[key] = entry
        self._attach_metrics(entry)
        return entry

    @contextmanager
    def lease(self, key, engine_factory):
        with self._lock:
            entry = self._entry_locked(key, engine_factory)
            entry.leases += 1
            entry.last_used_at = self._clock()
        try:
            yield entry.engine
        finally:
            with self._lock:
                entry.leases = max(0, entry.leases - 1)
                entry.last_used_at = self._clock()
                if entry.draining and entry.leases == 0:
                    self._dispose(entry)
                    if entry in self._draining:
                        self._draining.remove(entry)

    def invalidate(self, data_source_uid, reason):
        del reason  # The caller records the bounded audit reason.
        with self._lock:
            for entry in list(self._entries.values()):
                if entry.key.data_source_uid == data_source_uid:
                    self._start_draining_locked(entry)

    def expire_draining(self, now=None):
        now = self._clock() if now is None else float(now)
        disposed = 0
        with self._lock:
            for entry in list(self._draining):
                if now >= entry.drain_deadline:
                    self._dispose(entry)
                    self._draining.remove(entry)
                    disposed += 1
        return disposed

    def evict_idle(self, now=None):
        now = self._clock() if now is None else float(now)
        evicted = []
        with self._lock:
            idle = [
                entry
                for entry in self._entries.values()
                if entry.leases == 0 and not entry.draining
            ]
            ttl = float(self.settings.get("idle_ttl", 900))
            for entry in idle:
                if now - entry.last_used_at >= ttl:
                    evicted.append(entry)

            remaining = [
                entry for entry in idle if entry not in evicted
            ]
            maximum = int(self.settings.get("max_idle_pools", 20))
            overflow = max(0, len(remaining) - maximum)
            if overflow:
                evicted.extend(
                    sorted(remaining, key=lambda item: item.last_used_at)[
                        :overflow
                    ]
                )

            for entry in evicted:
                self._entries.pop(entry.key, None)
                self._dispose(entry)
        return len(evicted)

    def close_all(self):
        with self._lock:
            entries = list(self._entries.values()) + list(self._draining)
            self._entries.clear()
            self._draining.clear()
            for entry in entries:
                self._dispose(entry)

    def breaker_for(self, key):
        with self._lock:
            entry = self._entries.get(key)
            return entry.breaker if entry is not None else None

    def add_checkout_wait(self, key, elapsed_ms):
        with self._lock:
            entry = self._entries.get(key)
            if entry is not None:
                entry.checkout_wait_ms = float(elapsed_ms)

    def record_connection_failure(self, key):
        with self._lock:
            entry = self._entries.get(key)
            if entry is not None:
                entry.metrics["connection_failed_total"] += 1
                entry.breaker.record_failure()

    def record_connection_success(self, key):
        with self._lock:
            entry = self._entries.get(key)
            if entry is not None:
                entry.breaker.record_success()

    def record_pool_timeout(self, key):
        with self._lock:
            entry = self._entries.get(key)
            if entry is not None:
                entry.metrics["pool_timeout_total"] += 1

    def record_query_duration(self, key, elapsed_ms):
        with self._lock:
            entry = self._entries.get(key)
            if entry is not None:
                entry.metrics["query_total"] += 1
                entry.metrics["last_query_duration_ms"] = float(elapsed_ms)

    @staticmethod
    def _pool_value(engine, name):
        pool = getattr(engine, "pool", None)
        method = getattr(pool, name, None)
        if not callable(method):
            return 0
        try:
            return int(method())
        except Exception:
            return 0

    def _status(self, entry):
        breaker = entry.breaker.snapshot()
        if entry.draining:
            state = "draining"
        elif breaker["state"] == "open":
            state = "open"
        elif breaker["consecutive_failures"]:
            state = "degraded"
        else:
            state = "healthy"
        metrics = entry.metrics
        return PoolStatus(
            data_source_uid=entry.key.data_source_uid,
            credential_version=entry.key.credential_version,
            config_fingerprint=entry.key.config_fingerprint,
            pool_state=state,
            leases=entry.leases,
            draining=entry.draining,
            created_at=entry.created_at,
            last_used_at=entry.last_used_at,
            pool_size=self._pool_value(entry.engine, "size"),
            checked_out=self._pool_value(entry.engine, "checkedout"),
            checked_out_peak=metrics["checked_out_peak"],
            checked_in=self._pool_value(entry.engine, "checkedin"),
            overflow=self._pool_value(entry.engine, "overflow"),
            checkout_wait_ms=entry.checkout_wait_ms,
            connection_created_total=metrics["connection_created_total"],
            connection_failed_total=metrics["connection_failed_total"],
            pool_timeout_total=metrics["pool_timeout_total"],
            invalidated_total=metrics["invalidated_total"],
            query_total=metrics["query_total"],
            last_query_duration_ms=metrics["last_query_duration_ms"],
            consecutive_failures=breaker["consecutive_failures"],
            circuit_open_until=breaker["open_until"],
        )

    def snapshot(self):
        with self._lock:
            entries = list(self._entries.values()) + list(self._draining)
            return [self._status(entry) for entry in entries]
