import json

from loguru import logger

from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
from bb_integrations_lib.protocols.pipelines import Step
from bb_integrations_lib.shared.model import SDSetOrderExportStatusRequest, ERPStatus, SDGetUnexportedOrdersResponse


class BBDFinalizeExportStatuses(Step):
    """
    Final step for order export pipelines. Sets export statuses for all orders
    based on pipeline results. Reads all data from pipeline_context.extra_data:
        - unexported_orders: original orders list
        - errored_orders: orders that errored during processing
        - processed_orders: orders that were successfully processed
        - rollback_processed: if True, rolls back processed orders to 'pending'

    On rollback, processed orders are set to 'pending'. Otherwise, they are set to 'sent'.
    Any orders not accounted for in errors or processed are marked as 'errors'.
    """

    def __init__(self,
                 sd_client: GravitateSDAPI,
                 step_key: str | None = "Order Export Results",
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)
        self.sd_client = sd_client
        self.step_key = step_key

    def describe(self):
        return "Collect all results and set order export status"

    async def execute(self) -> None:
        orders: list[SDGetUnexportedOrdersResponse] = self.pipeline_context.extra_data.get("unexported_orders", [])
        errors = self.pipeline_context.extra_data.get("errored_orders", [])
        processed = self.pipeline_context.extra_data.get("processed_orders", [])
        rollback_processed = self.pipeline_context.extra_data.get("rollback_processed", False)

        all_requests: list[SDSetOrderExportStatusRequest] = []

        errored_orders_requests = format_status_requests(errors, ERPStatus.errors)
        all_requests.extend(errored_orders_requests)

        if rollback_processed:
            rollback_order_requests = format_status_requests(processed, ERPStatus.pending)
            all_requests.extend(rollback_order_requests)
            logger.warning("Error while exporting. Rolling back order status to 'pending'.")
            rollback_requests = json.dumps([order.model_dump() for order in rollback_order_requests], indent=2)
            self.pipeline_context.included_files[f"{self.step_key} - rolled back orders"] = rollback_requests
        else:
            processed_orders_requests = format_status_requests(processed, ERPStatus.sent)
            all_requests.extend(processed_orders_requests)
            processed_json = json.dumps([order.model_dump() for order in processed_orders_requests], indent=2)
            self.pipeline_context.included_files[f"{self.step_key} - processed orders"] = processed_json


        # Set filtered orders back to pending
        filtered = self.pipeline_context.extra_data.get("filtered_orders", [])
        if filtered:
            filtered_requests = [
                SDSetOrderExportStatusRequest(order_id=str(order.order_id), status=ERPStatus.pending)
                for order in filtered
            ]
            all_requests.extend(filtered_requests)

        # Catch any orders that were in the input but never ended up in errors or processed
        accounted_ids = {str(order.get("_id") or order.get("order_id")) for order in errors + processed}
        missing_order_ids = [ str(order.order_id) for order in orders if str(order.order_id) not in accounted_ids]
        missing_orders_requests = []
        if missing_order_ids:
            logger.error(f"These orders were missed by the pipeline: {missing_order_ids}")
            missing_orders_requests = [
                SDSetOrderExportStatusRequest(
                    order_id=str(order.order_id),
                    status=ERPStatus.errors,
                    error="Missed by the pipeline" #TODO: should this be errors
                )
                for order in orders
                if str(order.order_id) not in accounted_ids
            ]
            all_requests.extend(missing_orders_requests)

        errored_json = json.dumps([order.model_dump() for order in errored_orders_requests + missing_orders_requests], indent=2)
        self.pipeline_context.included_files[f"{self.step_key} - errored orders"] = errored_json

        # Send all requests at once
        response = await self.sd_client.bulk_set_export_order_status(all_requests)
        response.raise_for_status()

def format_status_requests(orders: list[dict], status: ERPStatus) -> list[SDSetOrderExportStatusRequest]:
    status_requests = [
        SDSetOrderExportStatusRequest(
            order_id=str(order.get("_id") or order.get("order_id")),
            status=status,
            error=order.get("error_message") if status == ERPStatus.errors else None,
        )
        for order in orders
    ]

    return status_requests
