import os
import pymongo
from pymongo import MongoClient
import urllib.parse


class MongoDB:
    def __init__(self, connection_string: str):
        self.client = MongoClient(connection_string)
        self.db = self.client[os.getenv("MONGO_DB")]  # use env DB name

    def close(self):
        self.client.close()


# MongoDB connection setup
def get_mongo_db():
    mongo_host = os.getenv("MONGO_HOST", "localhost")
    mongo_port = os.getenv("MONGO_PORT", "27017")
    mongo_user = os.getenv("MONGO_USER")
    mongo_pass = os.getenv("MONGO_PASS")
    mongo_db_name = os.getenv("MONGO_DB", "movex")
    mongo_uri = os.getenv("MONGO_URI")

    if not mongo_uri:
        if mongo_user and mongo_pass:
            # Ensure password is URL encoded if it contains special characters
            mongo_pass = urllib.parse.quote_plus(mongo_pass)
            mongo_uri = (
                f"mongodb://{mongo_user}:{mongo_pass}@{mongo_host}:{mongo_port}/{mongo_db_name}"
                f"?authSource={mongo_db_name}"
            )
        else:
            mongo_uri = f"mongodb://{mongo_host}:{mongo_port}/{mongo_db_name}"

    client = pymongo.MongoClient(mongo_uri)
    db = client[mongo_db_name]
    try:
        yield db
    finally:
        db.client.close()