from typing import Callable

from loguru import logger

from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
from bb_integrations_lib.models.pipeline_structs import PipelineContext, NoPipelineData
from bb_integrations_lib.pipelines.shared.allocation_matcher.matching_utils import match_allocations
from bb_integrations_lib.protocols.pipelines import Step
from bb_integrations_lib.shared.model import SDGetUnexportedOrdersResponse, GetOrderBolsAndDropsRequest


class BBDGetBolsAndDrops(Step):

    def __init__(self,
                 sd_client: GravitateSDAPI | None = None,
                 test_order_ids: list[str] | None = None,
                 include_allocation_matching: bool = False,
                 order_filter: Callable | None = None,
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)
        if self.pipeline_context is None:
            self.pipeline_context = PipelineContext()
        self.sd_client = sd_client
        self.test_override_order_ids = test_order_ids
        self.include_allocation_matching = include_allocation_matching
        self.order_filter = order_filter

        self.errors = []

    def describe(self):
        return "Get Bols and Drops from Supply and Dispatch. Optionally add allocation matching info to orders."

    async def execute(self, orders: list[SDGetUnexportedOrdersResponse] | None = None) -> list[dict]:
        if not orders:
            raise NoPipelineData

        if self.test_override_order_ids is not None:
            order_ids = self.test_override_order_ids
        else:
            order_ids = [order.order_id for order in orders]

        raw_bols_and_drops = await self.get_bols_and_drops(order_ids)

        order_bols_and_drops = self.catch_missing_and_errors(order_ids, raw_bols_and_drops)

        if self.include_allocation_matching:
            order_bols_and_drops = await self.add_allocation_matching_info(
                order_bols_and_drops
            )

        if self.order_filter:
            filtered_out_ids = {
                str(order.get("_id") or order.get("order_id"))
                for order in order_bols_and_drops
                if not self.order_filter(order)
            }
            unexported_orders = self.pipeline_context.extra_data.get("unexported_orders", [])
            for order in unexported_orders[:]:
                if order.order_id in filtered_out_ids:
                    unexported_orders.remove(order)
                    self.pipeline_context.extra_data.setdefault("filtered_orders", []).append(order)

            order_bols_and_drops = [order for order in order_bols_and_drops if self.order_filter(order)]

        # Flush all collected errors to pipeline context
        if self.errors:
            self.pipeline_context.extra_data.setdefault("errored_orders", []).extend(self.errors)

        return order_bols_and_drops

    async def get_bols_and_drops(self, order_ids: list[str]) -> list[dict]:
        if not order_ids:
            return []
        req = GetOrderBolsAndDropsRequest(order_ids=order_ids)
        response = await self.sd_client.get_bols_and_drops(req)
        return response.json()

    def catch_missing_and_errors(self, unexported_orders: list[str], orders: list[dict]) -> list[dict]:
        # Catch orders that are not returned by the get_bols_and_drops api
        if unexported_orders is not None:
            returned_ids = {str(order.get("_id") or order.get("order_id")) for order in orders}
            missing_bols_and_drops = [
                {"order_id": order_id, "error_message": "Could not retrieve bols_and_drops from the api"}
                for order_id in unexported_orders if order_id not in returned_ids
            ]
            self.errors.extend(missing_bols_and_drops)

        # Remove orders without allocated bols
        allocation_errors = [order for order in orders if order.get('allocated_bol_error')]
        for order in allocation_errors:
            order["error_message"] = order.get("allocated_bol_error")
        self.errors.extend(allocation_errors)

        orders = [order for order in orders if not order.get('allocated_bol_error')]

        return orders

    async def get_planned_orders(
            self,
            order_numbers: list[str | int],
    ) -> dict[str, dict | None]:
        planned_orders = await self.sd_client.get_orders_by_numbers(order_numbers, 7)

        # Fill in None for any requested order numbers not found
        for n in order_numbers:
            if str(n) not in planned_orders:
                planned_orders[str(n)] = None

        return planned_orders

    async def add_allocation_matching_info(
            self,
            order_bols_and_drops: list[dict],
    ) -> list[dict]:
        """
        For each order:
        1. Get planned order data
        2. Run allocation matching
        3. Add 'matched_allocations' to the order

        Args:
            order_bols_and_drops: list of bols_and_drops responses (one per order)

        Returns:
            The same list + 'matched_allocations' field
        """
        order_numbers = [
            str(bd.get("order_number"))
            for bd in order_bols_and_drops
            if bd.get("order_number")
        ]

        planned_orders = await self.get_planned_orders(order_numbers)

        matched_orders = []
        for order_data in order_bols_and_drops:
            order_num = str(order_data.get("order_number", ""))
            planned_order = planned_orders.get(order_num)

            try:
                if not planned_order:
                    order_data["error_message"] = f"No planned order found for order {order_num}"
                    self.errors.append(order_data)
                    continue
                order_data["planned_order"] = planned_order

                allocated_bols = order_data.get("allocated_bols", [])
                executed_bols = order_data.get("bols", [])
                executed_drops = order_data.get("drops", [])
                planned_loads = planned_order.get("loads", [])
                planned_drops = planned_order.get("drops", [])

                if not allocated_bols and not planned_drops:
                    order_data["matched_allocations"] = []
                    matched_orders.append(order_data)
                    continue

                matched = match_allocations(
                    order_number=order_num,
                    allocated_bols=allocated_bols,
                    executed_bols=executed_bols,
                    executed_drops=executed_drops,
                    planned_loads=planned_loads,
                    planned_drops=planned_drops,
                )

                order_data["matched_allocations"] = [
                    ma.model_dump() for ma in matched
                ]

                matched_orders.append(order_data)

            except Exception as e:
                msg = f"Allocation matching failed for order {order_num}: {e}"
                logger.warning(msg)
                order_data["error_message"] = msg
                self.errors.append(order_data)

        return matched_orders
