from __future__ import annotations

from dataclasses import dataclass
from threading import RLock
from typing import Any

from motor.motor_asyncio import AsyncIOMotorClient
from pymongo import MongoClient

from app.config import config
from app.core.logging import InstrumentedAsyncMongoDatabase


@dataclass
class MongoRuntime:
    sync_client: MongoClient
    async_client: AsyncIOMotorClient
    sync_database: Any
    async_database: InstrumentedAsyncMongoDatabase

    @classmethod
    def create(
        cls,
        *,
        mongo_uri: str | None = None,
        db_name: str | None = None,
    ) -> MongoRuntime:
        normalized_uri = str(mongo_uri or config.MONGO_URI).strip()
        normalized_db_name = str(db_name or config.DB_NAME).strip()
        sync_client = MongoClient(normalized_uri)
        async_client = AsyncIOMotorClient(normalized_uri)
        sync_database = sync_client[normalized_db_name]
        async_database = InstrumentedAsyncMongoDatabase(
            async_client[normalized_db_name],
            service="mongo_async",
        )
        return cls(
            sync_client=sync_client,
            async_client=async_client,
            sync_database=sync_database,
            async_database=async_database,
        )

    async def ping(self) -> None:
        await self.async_client.admin.command("ping")

    def close(self) -> None:
        self.async_client.close()
        self.sync_client.close()


_runtime_lock = RLock()
_current_runtime: MongoRuntime | None = None


def set_mongo_runtime(runtime: MongoRuntime | None) -> None:
    global _current_runtime
    with _runtime_lock:
        _current_runtime = runtime


def init_mongo_runtime(
    *,
    mongo_uri: str | None = None,
    db_name: str | None = None,
) -> MongoRuntime:
    runtime = MongoRuntime.create(mongo_uri=mongo_uri, db_name=db_name)
    previous = None
    with _runtime_lock:
        global _current_runtime
        previous = _current_runtime
        _current_runtime = runtime
    if previous is not None and previous is not runtime:
        previous.close()
    return runtime


def get_mongo_runtime() -> MongoRuntime:
    global _current_runtime
    with _runtime_lock:
        if _current_runtime is None:
            _current_runtime = MongoRuntime.create()
        return _current_runtime


def clear_mongo_runtime(runtime: MongoRuntime | None = None, *, close: bool = True) -> None:
    global _current_runtime
    with _runtime_lock:
        current = _current_runtime
        if runtime is not None and current is not runtime:
            if close:
                runtime.close()
            return
        _current_runtime = None
    if current is not None and close:
        current.close()