from db import db

from workers.base_worker import BaseWorker
from workers.prs.repository import PRSImportRepository
from workers.prs.mappers.budget_mapper import BudgetMapper
from workers.prs.models import BudgetModel


class PRSBudgetWorker(BaseWorker):
    repository = PRSImportRepository()
    worker_name = "PRSBudgetWorker"

    def should_log_exceptions(self):
        return True

    def execute(self, *args):
        budgets = self.repository.get_raw_prs_budgets()
        self.logger.info("PRS Budgets import", f"Importing {len(budgets)} budgets...")

        for result in budgets:
            project_id = result.project_id

            mapped_budget = BudgetMapper(
                raw_budget=BudgetModel(
                    project_id=project_id,
                    group_id=result.group_id,
                    category_id=result.category_id,
                    value=result.proposed_am,
                    prs_project_code=result.sub_project_cd,
                    external_key=self.__external_key(result),
                )
            ).build()
            db.session.add(mapped_budget)
            db.session.flush()

        self.logger.success("PRS Budgets import", "Finished!")

        db.session.commit()
        db.session.close()

    def __external_key(self, budget):
        return f"""{budget.sub_project_cd}-{budget.form_id}-{budget.phase_id}-{budget.group_id}-{
                    budget.category_id}-{budget.created_at}"""