from models.prs_project_budget import PRSProjectBudget
from workers.prs.repository import PRSImportRepository


class BudgetMapper:
    repository = PRSImportRepository()

    def __init__(self, raw_budget):
        self.raw_budget = raw_budget

    def build(self):
        budget = self.__load_budget_entity() or PRSProjectBudget()
        budget.project_id = self.raw_budget.project_id
        budget.group_id = self.raw_budget.group_id
        budget.category_id = self.raw_budget.category_id
        budget.value = self.raw_budget.value
        budget.external_key = self.raw_budget.external_key
        budget.currency = self.raw_budget.currency
        return budget

    def __load_budget_entity(self):
        if self.raw_budget.external_key:
            return self.repository.get_budget_by_external_key(self.raw_budget.external_key)
        return None