from datetime import date, datetime, timedelta, UTC
from bb_integrations_lib.protocols.pipelines import Step
from bb_integrations_lib.provider.api.samsara.client import SamsaraClient
from loguru import logger


class SamsaraRetrieveHOSDataStep(Step):
    """Retrieve HOS clocks and daily logs from Samsara for each driver."""

    def __init__(self,
                 samsara_client: SamsaraClient,
                 hos_clock_hours_back: int = 24,
                 *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.samsara_client = samsara_client
        self.hos_clock_hours_back = hos_clock_hours_back

    def describe(self) -> str:
        return "Retrieve HOS data from Samsara API"

    @staticmethod
    def logs_lookup_by_driver(logs_data: list[dict]) -> dict[str, dict]:
        """Build lookup from driver ID to daily log, taking first log per driver."""
        lookup = {}
        for log in logs_data:
            driver_id = log["driver"]["id"]
            if driver_id not in lookup:
                lookup[driver_id] = log
        return lookup

    async def retrieve_clocks_and_daily_logs(
            self,
            samsara_ids: list[str],
            today: date,
    ) -> tuple[dict[str, dict], dict[str, dict]]:
        """Retrieve HOS clocks and daily logs from Samsara."""
        logger.info(f"Retrieving HOS clocks for {len(samsara_ids)} drivers")
        clocks_response = await self.samsara_client.get_hos_clocks(samsara_ids)
        clocks_data = clocks_response.json().get("data", [])
        clocks_by_driver = {
            c["driver"]["id"]: c.get("clocks", {}) for c in clocks_data
        }

        logger.info(f"Retrieving HOS daily logs for {len(samsara_ids)} drivers")
        logs_response = await self.samsara_client.get_hos_daily_logs(
            driver_ids=samsara_ids, start_date=today, end_date=today
        )
        logs_data = logs_response.json().get("data", [])
        logs_by_driver = self.logs_lookup_by_driver(logs_data)

        return clocks_by_driver, logs_by_driver

    async def retrieve_hos_log_events_by_driver(
            self,
            samsara_ids: list[str],
            start_time: datetime,
            end_time: datetime,
    ) -> dict[str, list]:
        """Retrieve HOS log segments (duty status events) from Samsara."""
        logger.info(f"Retrieving HOS log segments for {len(samsara_ids)} drivers")
        hos_logs_response = await self.samsara_client.get_hos_logs(
            driver_ids=samsara_ids,
            start_time=start_time,
            end_time=end_time,
        )
        hos_logs_data = hos_logs_response.json().get("data", [])
        return {
            entry["driver"]["id"]: entry.get("hosLogs", []) for entry in hos_logs_data
        }

    @staticmethod
    def combine_driver_data(
            samsara_ids: list[str],
            samsara_to_gravitate: dict[str, dict],
            clocks_by_driver: dict[str, dict],
            logs_by_driver: dict[str, dict],
            events_by_driver: dict[str, list],
    ) -> list[dict]:
        """Combine all retrieved data into result records for each driver."""
        results = []
        retrieve_timestamp = datetime.now(UTC)

        for samsara_id in samsara_ids:
            gravitate_driver = samsara_to_gravitate.get(samsara_id)
            if not gravitate_driver:
                continue

            results.append({
                "gravitate_driver": gravitate_driver,
                "samsara_id": samsara_id,
                "clocks": clocks_by_driver.get(samsara_id, {}),
                "daily_log": logs_by_driver.get(samsara_id, {}),
                "hos_events": events_by_driver.get(samsara_id, []),
                "retrieve_timestamp": retrieve_timestamp,
            })

        return results

    async def execute(self, drivers: list[dict]) -> list[dict]:
        """
        Retrieve HOS data from Samsara for all drivers.

        Args:
            drivers: List of Gravitate drivers with samsara_id in extra_data

        Returns:
            List of dicts containing driver info + Samsara HOS data
        """
        logs_start = datetime.now(UTC) - timedelta(hours=self.hos_clock_hours_back)
        logs_end = datetime.now(UTC)
        today = date.today()

        if not drivers:
            logger.info("No drivers with Samsara IDs found")
            return []

        # Extract samsara_ids
        samsara_ids = [d["extra_data"]["samsara_id"] for d in drivers]

        # Build lookup from samsara_id to gravitate driver
        samsara_to_gravitate = {d["extra_data"]["samsara_id"]: d for d in drivers}

        clocks_by_driver, logs_by_driver = await self.retrieve_clocks_and_daily_logs(
            samsara_ids, today
        )
        events_by_driver = await self.retrieve_hos_log_events_by_driver(
            samsara_ids, logs_start, logs_end
        )
        results = self.combine_driver_data(
            samsara_ids,
            samsara_to_gravitate,
            clocks_by_driver,
            logs_by_driver,
            events_by_driver,
        )
        logger.info(f"Retrieved HOS data for {len(results)} drivers")
        return results

