# core/database.py

from pymongo import MongoClient
from core.files import Data

# Initialize the client once at the module level
client = MongoClient(Data('config').yaml_read()['mongo'])

class Servers:
    def __init__(self, server_id=None):
        self.server_id = server_id
        # All server-related data now lives in a single collection
        self.collection = client["BytesBump"]["servers"]

    def get(self):
        """Finds a single server document by its ID."""
        return self.collection.find_one({"_id": self.server_id})

    def get_all(self):
        """Returns a list of all documents in the servers collection."""
        return list(self.collection.find({}))

    def add(self, **params):
        """Adds a new server document to the database."""
        params['_id'] = self.server_id
        self.collection.insert_one(params)

    def update(self, **params):
        """Updates fields in a server document."""
        self.collection.update_one(
            {"_id": self.server_id},
            {"$set": params}
        )

    def delete(self):
        """Removes a server document from the database."""
        if self.server_id:
            self.collection.delete_one({'_id': self.server_id})

    # --- REFACTORED PREFIX MANAGEMENT ---

    def get_prefix(self):
        """Gets a custom prefix for a server from its document."""
        server_data = self.get()
        # Return the prefix if it exists in the document, otherwise return None
        if server_data:
            return server_data.get('prefix')
        return None

    def set_prefix(self, prefix: str):
        """Sets or updates the prefix for a server."""
        # This will add the 'prefix' field to the document or update its value.
        self.update(prefix=prefix)

    def delete_prefix(self):
        """Removes the custom prefix from a server document."""
        # The $unset operator is the correct way to remove a field from a document.
        self.collection.update_one(
            {"_id": self.server_id},
            {"$unset": {"prefix": ""}}
        )