from __future__ import annotations

import os
import threading
import weakref
from typing import Any, Dict, Optional, Type

from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    PrivateAttr,
    SecretStr,
    model_validator,
)
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.engine import url as sqlalchemy_url
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import NullPool, Pool, QueuePool, StaticPool

from sibi_flux.logger import Logger

# Ensure this import exists or mock it if strictly testing this file
from ._sql_model_builder import SqlAlchemyModelBuilder

# --- Global Engine Registry ---
# Stores {key: {"engine": Engine, "ref_count": int}}
_ENGINE_REGISTRY_LOCK = threading.RLock()
_ENGINE_REGISTRY: Dict[tuple, Any] = {}


def _release_engine_resource(key: tuple, logger: Optional[Logger] = None) -> None:
    """
    Standalone clean-up function for weakref.finalize.
    Decrements ref_count and disposes the engine if count reaches 0.
    """
    # Guard against interpreter shutdown (globals might be None)
    if _ENGINE_REGISTRY_LOCK is None or _ENGINE_REGISTRY is None:
        return

    try:
        with _ENGINE_REGISTRY_LOCK:
            if key in _ENGINE_REGISTRY:
                wrapper = _ENGINE_REGISTRY[key]
                wrapper["ref_count"] -= 1

                if logger:
                    # Swallow logging errors during shutdown
                    try:
                        logger.debug(
                            f"DB Engine ref count: {wrapper['ref_count']} for {key}"
                        )
                    except Exception:
                        pass

                if wrapper["ref_count"] <= 0:
                    try:
                        wrapper["engine"].dispose()
                        if logger:
                            try:
                                logger.debug(f"Disposed DB Engine for key {key}")
                            except Exception:
                                pass
                    except Exception:
                        pass
                    finally:
                        _ENGINE_REGISTRY.pop(key, None)
    except Exception:
        # Final safety net to prevent errors during GC
        pass


class SqlAlchemyConnectionConfig(BaseModel):
    """
    Immutable, thread-safe SQLAlchemy connection manager.
    """

    model_config = ConfigDict(
        frozen=True,
        arbitrary_types_allowed=True,
        extra="forbid",
    )

    connection_url: SecretStr = Field(..., description="Database DSN")
    table: Optional[str] = None
    debug: bool = False
    logger_extra: Optional[Dict[str, Any]] = Field(
        default_factory=lambda: {"sibi_flux_component": __name__}
    )

    # --- Pool Configuration ---
    pool_size: int = Field(
        default_factory=lambda: int(os.environ.get("DB_POOL_SIZE", 5)), ge=0
    )
    max_overflow: int = Field(
        default_factory=lambda: int(os.environ.get("DB_MAX_OVERFLOW", 10)), ge=0
    )
    pool_timeout: int = Field(
        default_factory=lambda: int(os.environ.get("DB_POOL_TIMEOUT", 30)), ge=0
    )
    pool_recycle: int = Field(
        default_factory=lambda: int(os.environ.get("DB_POOL_RECYCLE", 1800))
    )
    pool_pre_ping: bool = Field(default=True)
    poolclass: Type[Pool] = QueuePool
    connect_args: Optional[Dict[str, Any]] = Field(
        default_factory=dict,
        description="Driver-specific connection arguments (e.g. timeouts)",
    )
    execution_options: Optional[Dict[str, Any]] = Field(
        default_factory=dict, description="Engine execution options"
    )

    # --- Runtime State ---
    _model: Optional[Any] = PrivateAttr(default=None)
    _engine: Optional[Engine] = PrivateAttr(default=None)
    _logger: Optional[Logger] = PrivateAttr(default=None)
    _session_factory: Optional[sessionmaker] = PrivateAttr(default=None)
    _engine_key_instance: Optional[tuple] = PrivateAttr(default=None)
    _finalizer: Any = PrivateAttr(default=None)

    @model_validator(mode="after")
    def _initialize_resources(self) -> "SqlAlchemyConnectionConfig":
        if self._logger is None:
            logger = Logger.default_logger(logger_name=self.__class__.__name__)
            logger.set_level(Logger.DEBUG if self.debug else Logger.INFO)
            self._logger = logger

        self._engine_key_instance = self._get_engine_key()
        self._init_engine()

        if self.table and self._engine:
            self._build_model()

        if self._engine:
            self._session_factory = sessionmaker(
                bind=self._engine, expire_on_commit=False
            )

        return self

    def _get_engine_key(self) -> tuple:
        url_str = self.connection_url.get_secret_value()
        parsed = sqlalchemy_url.make_url(url_str)

        # Remove volatile query params if necessary, usually pool params aren't in query
        query = {k: v for k, v in parsed.query.items()}
        normalized_url = parsed.set(query=query)

        key_parts = [str(normalized_url), self.poolclass]

        # Only include pool size params if the pool supports them
        if self.poolclass not in (NullPool, StaticPool):
            key_parts.extend(
                [
                    self.pool_size,
                    self.max_overflow,
                    self.pool_timeout,
                ]
            )

        # All pools support recycle/pre_ping logic in create_engine wrapper
        key_parts.extend([self.pool_recycle, self.pool_pre_ping])

        return tuple(key_parts)

    def _init_engine(self) -> None:
        key = self._engine_key_instance
        if key is None:
            raise ValueError("Engine key is None")

        with _ENGINE_REGISTRY_LOCK:
            wrapper = _ENGINE_REGISTRY.get(key)
            if wrapper:
                self._engine = wrapper["engine"]
                wrapper["ref_count"] += 1
                if self.debug:
                    if self._logger:
                        self._logger.debug(
                            f"Reusing DB engine. Ref count: {wrapper['ref_count']}"
                        )
            else:
                if self.debug:
                    db_url = self.connection_url.get_secret_value()
                    safe_url = sqlalchemy_url.make_url(db_url).render_as_string(
                        hide_password=True
                    )
                    if self._logger:
                        self._logger.debug(f"Creating engine with url: {safe_url}")
                try:
                    # [FIX] Filter arguments based on poolclass
                    kwargs = {
                        "pool_recycle": self.pool_recycle,
                        "pool_pre_ping": self.pool_pre_ping,
                        "poolclass": self.poolclass,
                    }

                    if self.poolclass not in (NullPool, StaticPool):
                        kwargs.update(
                            {
                                "pool_size": self.pool_size,
                                "max_overflow": self.max_overflow,
                                "pool_timeout": self.pool_timeout,
                            }
                        )

                    new_engine = create_engine(
                        self.connection_url.get_secret_value(),
                        connect_args=self.connect_args,
                        execution_options=self.execution_options,
                        **kwargs,
                    )

                    _ENGINE_REGISTRY[key] = {
                        "engine": new_engine,
                        "ref_count": 1,
                    }
                    self._engine = new_engine
                except Exception as e:
                    if self._logger:
                        self._logger.error(f"Failed to create engine: {e}")
                    raise SQLAlchemyError(f"DB Engine creation failed: {e}") from e

        self._finalizer = weakref.finalize(
            self, _release_engine_resource, key, self._logger
        )

    def _build_model(self) -> None:
        try:
            if not self.table:
                raise ValueError("Table name is required")
            builder = SqlAlchemyModelBuilder(self._engine, self.table)
            self._model = builder.build_model()
        except Exception as e:
            if self._logger:
                self._logger.error(f"Failed to build ORM model: {e}")
            raise ValueError(f"Model construction failed: {e}") from e

    @property
    def engine(self) -> Engine:
        if self._engine is None:
            raise RuntimeError("Engine not initialized")
        return self._engine

    @property
    def model(self) -> Any:
        if self._model is None and self.table:
            # Lazy load fallback
            self._build_model()
        return self._model

    def close(self) -> None:
        if self._finalizer:
            self._finalizer()
