from News.utils.mongodb.mongo import MongoConnector
from Unternehmensregister.models.Company import Company
from Unternehmensregister.utils.CompanyServiceInterface import CompanyServiceInterface


class CompanyMongoService(CompanyServiceInterface):
    def __init__(self, connector: MongoConnector):
        self.collection = connector.database["companies"]

    def get_all(self) -> list[Company]:
        result = self.collection.find()
        return list(result)

    def get_by_id(self, id: str) -> Company | None:
        result = list(self.collection.find({"id": id}))
        if len(result) == 1:
            return result[0]
        return None

    def insert(self, company: Company):
        return self.collection.insert_one(company.dict())