from datetime import datetime, UTC

from dateutil.parser import parse
from loguru import logger

from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
from bb_integrations_lib.protocols.pipelines import Step, ParserBase
from bb_integrations_lib.shared.model import RawData


class ExportDriverPayStep(Step):
    def __init__(self,
                 sd_client: GravitateSDAPI,
                 parser: type[ParserBase],
                 parser_kwargs: dict | None = None,
                 payroll_status: str | None = None,
                 target_date: datetime | None = None,
                 *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.sd_client = sd_client
        self.custom_parser = parser
        self.custom_parser_kwargs = parser_kwargs or {}
        self.payroll_status = payroll_status
        self.target_date = target_date

    def describe(self):
        return "Export Driver Pay"

    async def execute(self, i=None) -> RawData:
        dt = self.right_date()
        target_date = dt.isoformat()
        params = {"date": target_date}
        if self.payroll_status:
            params["status"] = self.payroll_status

        logger.info(f"Fetching payroll export with params: {params}")
        resp = await self.sd_client.payroll_export(**params)
        data = resp.json()
        logger.info(f"Received {len(data)} payroll records")

        parser = self.custom_parser(**self.custom_parser_kwargs)
        return await parser.parse(data)

    def right_date(self) -> datetime:
        if self.target_date is not None:
            if isinstance(self.target_date, datetime):
                return self.target_date
            elif isinstance(self.target_date, str):
                return parse(self.target_date)
            else:
                raise RuntimeError(f"Target date {self.target_date} is not in a valid format")
        else:
            return datetime.now(UTC)
