import json

from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
from bb_integrations_lib.models.pipeline_structs import NoPipelineData
from bb_integrations_lib.protocols.pipelines import Step
from bb_integrations_lib.pipelines.steps.samsara.samsara_filter_changed_step import FilterResult
from loguru import logger
from pymongo.asynchronous.database import AsyncDatabase


class SamsaraUploadELDStep(Step):
    """Upload ELD data to Gravitate and update MongoDB hashes."""
    def __init__(
            self,
            sd_client: GravitateSDAPI,
            mongo_db: AsyncDatabase | None = None,
            collection_name: str = "HosSync",
            *args,
            **kwargs,
    ):
        self.sd_client = sd_client
        self.mongo_db = mongo_db
        self.collection_name = collection_name
        self.logs = []
        super().__init__(*args, **kwargs)

    def describe(self) -> str:
        return "Upload ELD data to Gravitate"

    async def update_hashes(self, hash_updates: list[dict]) -> None:
        """Update hash documents in MongoDB after successful upload."""
        if self.mongo_db is None or not hash_updates:
            return

        collection = self.mongo_db[self.collection_name]

        for update in hash_updates:
            await collection.update_one(
                {"driver_id": update["driver_id"]},
                {
                    "$set": {
                        "source_id": update["source_id"],
                        "hash": update["hash"],
                        "version": update["version"],
                        "updated_at": update["updated_at"],
                    }
                },
                upsert=True,
            )

        logger.info(f"Updated {len(hash_updates)} hash records in MongoDB")

    async def execute(self, filter_result: FilterResult) -> dict:
        """
        Upload ELD records to Gravitate and update MongoDB hashes.

        Args:
            filter_result: FilterResult from SamsaraFilterChangedStep containing
                          changed_records and hash_updates

        Returns:
            Dict with upload result info
        """
        eld_records = filter_result.changed_records
        hash_updates = filter_result.hash_updates

        if not eld_records:
            msg = "No ELD records to upload"
            logger.error(msg)
            raise NoPipelineData(msg)

        logger.info(f"Uploading {len(eld_records)} ELD records to Gravitate")

        response = await self.sd_client.upsert_eld_data(eld_records)
        result = response.json()

        logger.success(f"Uploaded {len(eld_records)} ELD records to Gravitate")

        await self.update_hashes(hash_updates)

        for record in eld_records:
            self.logs.append(
                {
                    "driver_id": record.get("driver_id"),
                    "source_id": record.get("source_id"),
                }
            )
        self.pipeline_context.included_files['SD request'] = json.dumps(eld_records)
        self.pipeline_context.included_files['Updated Driver HOS'] = json.dumps(self.logs)
        self.pipeline_context.included_files['SD response'] = json.dumps(result)
        return {"uploaded": len(eld_records), "response": result}
