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
