import json
from typing import Any

from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
from bb_integrations_lib.protocols.pipelines import Step
from loguru import logger


class SamsaraRetrieveDriversStep(Step):
    """Retrieve drivers from Gravitate
        Filter those with samsara_id in extra_data."""

    def __init__(self,
                 sd_client: GravitateSDAPI,
                 samsara_driver_extra_data_field: str = "samsara_id",
                 samsara_ids: list[str] | None = None,
                 *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.sd_client = sd_client
        self.samsara_driver_extra_data_field = samsara_driver_extra_data_field
        self.samsara_ids = samsara_ids

    def describe(self) -> str:
        return "Fetch drivers with Samsara ID from Gravitate"

    async def execute(self, data: Any) -> list[dict]:
        response = await self.sd_client.all_drivers()
        response.raise_for_status()
        drivers = response.json()

        # Filter drivers that have samsara_id in extra_data
        samsara_drivers = [
            driver
            for driver in drivers
            if driver.get("extra_data", {}).get(self.samsara_driver_extra_data_field, None)
        ]

        # Filter to specific Samsara IDs if provided
        if self.samsara_ids:
            samsara_drivers = [
                driver
                for driver in samsara_drivers
                if driver.get("extra_data", {}).get(self.samsara_driver_extra_data_field) in self.samsara_ids
            ]
            logger.info(f"Filtered to {len(samsara_drivers)} drivers matching requested Samsara IDs")

        logger.info(
            f"Found {len(samsara_drivers)} drivers with Samsara IDs out of {len(drivers)} total"
        )

        self.pipeline_context.included_files['S&D Samsara Drivers'] = json.dumps(samsara_drivers)
        return samsara_drivers