import time
from datetime import UTC, date, datetime
from typing import AsyncIterator, Iterable, Iterator

from loguru import logger
from path_dict import PathDict, pd

from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
from bb_integrations_lib.mappers.rita_mapper import RitaMapper, RitaMapperKeyError
from bb_integrations_lib.models.fuelquest import (
    DeliveryReport,
    FuelQuestNAXMLDelivery,
    TransmissionHeader,
)
from bb_integrations_lib.models.fuelquest.delivery import (
    BillTo,
    Consignor,
    DeliveredTo,
    DeliveryConfirmation,
    DeliveryEquipment,
    DeliveryInformation,
    LineItem,
    LineItemAllowanceOrCharge,
    LineItemFuelProductQty,
    LineItemTank,
    LineItemTankFuelProductId,
    LineItemTankMeasurement,
)
from bb_integrations_lib.models.fuelquest.invoice import (
    AdditionalAllowancesAndCharges,
    FreightCharge,
    FreightDetail,
    FreightInvoiceDetail,
    FreightInvoiceSummary,
    FreightInvoiceTotals,
    FreightInvTotalsDetails,
    FuelFreightInvoice,
    FuelProductId,
    FuelProductQty,
    InvoiceAllowanceOrCharge,
    InvoiceLocation,
    InvoiceSupplier,
    ProductDetail,
    RateDistance,
    RateReference,
    Terms,
    TotalFreightInvoiceDueAmt,
)
from bb_integrations_lib.models.fuelquest.shared import (
    AllowanceOrChargeReason,
    AllowanceOrChargeReasonValues,
    Buyer,
    Carrier,
    CountryCode,
    OrganizationId,
    ReleaseNumber,
    ReportType,
    SupplyTerminal,
    TruckTerminal,
    UOMBasis,
)
from bb_integrations_lib.models.rita.issue import IssueBase, IssueCategory
from bb_integrations_lib.models.rita.mapping import CompositeMapKey, MappingType


class TimedLog(object):
    def __init__(self, level: str, message: str):
        self.level = level
        self.message = message
        self.start = None

    def __enter__(self):
        self.start = time.monotonic()

    def __exit__(self, exc_type, exc_val, exc_tb):
        end = time.monotonic()
        logger.log(self.level, f"{self.message} (took {end - self.start:.3f}s)")  # ty:ignore[unsupported-operator]


class FuelQuestDeliveryPlanConverter:
    def __init__(
        self,
        rita_mapper: RitaMapper,
        sd_client: GravitateSDAPI,
        transmission_datetime: datetime,
        self_carrier: Carrier,
        transmission_sender: str,
        transmission_receiver: str,
        allowed_counterparties: list[str],
        allowed_store_markets: list[str] | None,
        country_code: str = "US",
        currency: str = "USD",
    ):
        self.rita_mapper = rita_mapper
        self.sd_client = sd_client
        self.transmission_datetime = transmission_datetime
        self.self_carrier = self_carrier
        self.transmission_sender = transmission_sender
        self.transmission_receiver = transmission_receiver
        self.allowed_counterparties = allowed_counterparties
        self.allowed_store_markets = allowed_store_markets
        self.country_code = country_code
        self.currency = currency

        self.report_type = ReportType.ORIGINAL  # TODO: Support UPDATEs?
        self.driver_sched_lkp = pd({})
        self.issues: list[IssueBase] = []
        self.skipped_orders_filter: list[tuple[str, str]] = []
        self.successful_order_numbers: list[str] = []

    def _buyer_from_allocated_bol(self, bol: PathDict) -> str:
        # TODO: Supply owner? We probably shouldn't use the store counterparty for buyer (but supply owner is not
        #   available on bols_and_drops)
        store_id = bol["store_id"]
        store = self.stores_by_id[store_id] or {}
        store_cp = self.cp_by_id[store["counterparty_id"]] or {}
        return store_cp["name"]

    def _rmke_to_issue(
        self, e: RitaMapperKeyError, order_number: int, extra_context: str | None = None
    ) -> IssueBase:
        logger.error(f"Error processing order {order_number}: {e}")
        if isinstance(e.from_key, CompositeMapKey):
            key_shortened = ",".join(
                f"{key}='{val}'" for key, val in e.from_key.key.items()
            )
        else:
            key_shortened = repr(e.from_key)
        context_str = f": {extra_context}" if extra_context else ""
        return IssueBase(
            key=f"mapping_failed_{key_shortened}",
            config_id="",
            name="Failed to map an item",
            category=IssueCategory.MISC,  # We don't have a "mapping" category yet
            problem_short=(
                f"Failed to map an item for order {order_number}{context_str}"
            ),
            problem_long=str(e),
        )  # ty:ignore[missing-argument]

    def _build_freight_invoice_detail(
        self,
        item: PathDict,
        invoices: Iterable[dict],
        order_number: str,
        terminal: SupplyTerminal,
        store: dict,
        store_fq_id: str,
        supplier_cp_id: str,
        fq_product_id: str,
    ) -> FreightInvoiceDetail:
        bol_number = item["bol_number"]
        delivered_qty = item["bol_net_volume_allocated"]

        # Find matching freight charge data from invoices (if any)
        charge_amt = 0.0
        rate = 0.0
        reason = AllowanceOrChargeReasonValues.FREIGHT
        for invoice in invoices:
            for transaction in invoice.get("transactions", []):
                if str(transaction["order_number"]) != order_number:
                    continue
                for detail in transaction.get("details", []):
                    if str(detail["bol_number"]) != str(bol_number):
                        continue
                    charge_amt = detail["total"]
                    rate = detail["rate"]
                    reason = {
                        "Base Freight": AllowanceOrChargeReasonValues.FREIGHT,
                        "Surcharge": AllowanceOrChargeReasonValues.FUEL_SURCHARGE,
                    }.get(detail["type"], AllowanceOrChargeReasonValues.OTHER_CHARGES)
                    break

        freight_detail = FreightDetail(
            count=1,
            allowance_or_charge=InvoiceAllowanceOrCharge(
                charge_amt=charge_amt,
                reason=AllowanceOrChargeReason(quantity=delivered_qty, reason=reason),
            ),
            freight_charge=FreightCharge(
                delivered_qty=delivered_qty,
                rated_qty=delivered_qty,
                rate=rate,
                rate_basis="",
                currency=self.currency,
                amount=charge_amt,
            ),
            product_details=[
                ProductDetail(
                    fuel_product_id=FuelProductId(
                        ident=fq_product_id, name=item["bol_product"]
                    ),
                    fuel_product_qty=FuelProductQty(
                        uom_basis=UOMBasis.NET_GALLONS, quantity=delivered_qty
                    ),
                )
            ],
        )

        # loading_dt = datetime.fromisoformat(bol_date)
        # delivery_dt = datetime.fromisoformat(delivered_date)
        loading_dt = datetime.now(UTC)
        delivery_dt = datetime.now(UTC)
        release_dt = datetime(1970, 1, 1, 0, 0, tzinfo=UTC)

        supplier_fq_id = self.rita_mapper.get_source_parent_id(
            self.cp_by_id[supplier_cp_id]["name"], MappingType.counterparty
        )

        return FreightInvoiceDetail(
            location=InvoiceLocation(
                name=store["store_number"],
                state=store["state"],
                country_code=CountryCode(code=self.country_code),
                organization_id=OrganizationId(ident=store_fq_id),
            ),
            supplier=InvoiceSupplier(
                name=self.cp_by_id[supplier_cp_id].get("name"),
                organization_id=OrganizationId(ident=supplier_fq_id),
            ),
            terminal=terminal,
            release_number=ReleaseNumber(
                expiration_datetime=release_dt,
                value="",
            ),
            bill_of_lading_number=str(bol_number),
            carrier_way_bill_number="",
            rate_reference=RateReference(
                ident=""  # Allowed to be blank
            ),
            rate_distance=RateDistance(uom_basis=UOMBasis.BLANK),
            loading_datetime=loading_dt,
            delivery_datetime=delivery_dt,
            freight_details=[freight_detail],
        )

    def _build_delivery_line_item(
        self, item: PathDict, fq_product_id: str, count: int
    ) -> LineItem:
        fq_tank_id = self.rita_mapper.get_source_child_id(
            item["store_number"], str(item["store_tank"]), MappingType.tank
        )
        return LineItem(
            count=count,
            tank=LineItemTank(ident=fq_tank_id),
            fuel_product_id=LineItemTankFuelProductId(
                ident=fq_product_id, value=item["bol_product"]
            ),
            fuel_product_qty=LineItemFuelProductQty(
                uom_basis=UOMBasis.NET_GALLONS, value=item["bol_net_volume_allocated"]
            ),
            # TODO: Do we need to implement tank measurements? Can we?
            tank_measurement=LineItemTankMeasurement(),
            allowance_or_charge=LineItemAllowanceOrCharge(
                charge_amt=0.0,
                allowance_or_charge_reason=AllowanceOrChargeReason(
                    quantity=0.0, reason=AllowanceOrChargeReasonValues.OTHER_CHARGES
                ),
            ),
        )

    def _build_confirmation_and_invoice_detail(
        self,
        order: PathDict,
        item: PathDict,
        invoices: Iterable[dict],
        po_number: str | None,
        truck_terminal: TruckTerminal,
        supply_terminal: SupplyTerminal,
        item_count: int,
    ) -> tuple[DeliveryConfirmation, FreightInvoiceDetail] | None:
        buyer_name = self._buyer_from_allocated_bol(item)
        if buyer_name not in self.allowed_counterparties:
            logger.debug(
                f"Skipping drop for non-enabled buyer counterparty {buyer_name}"
            )
            return None
        try:
            buyer_fq_id = self.rita_mapper.get_source_parent_id(
                buyer_name, MappingType.counterparty
            )
        except RitaMapperKeyError as e:
            # This can occur if there is an enabled but unmapped buyer counterparty
            logger.warning(
                f"No counterparty mapping found for enabled buyer '{buyer_name}' on order {order['order_number']}"
            )
            self.issues.append(
                self._rmke_to_issue(
                    e, order["order_number"], "(buyer counterparty to FQ ident)"
                )
            )
            return None

        store_id = item["store_id"]
        store = self.stores_by_id[store_id] or {}
        try:
            store_fq_id = self.rita_mapper.get_source_parent_id(
                store["store_number"], MappingType.site
            )
        except RitaMapperKeyError as e:
            logger.warning(e)
            self.issues.append(
                self._rmke_to_issue(
                    e, order["order_number"], "(store number to FQ ident)"
                )
            )
            return None

        buyer = Buyer(
            name=buyer_name,
            state=store["state"],
            country_code=CountryCode(code=self.country_code),
            organization_id=OrganizationId(
                ident=buyer_fq_id,
            ),
        )

        this_carrier = Carrier.model_validate(self.self_carrier.model_dump())
        this_carrier.truck_terminal = truck_terminal

        fq_product_id = self.rita_mapper.get_source_id_by_composite(
            CompositeMapKey(
                key={
                    "supplier": item["bol_supplier"],
                    "product name": item["bol_product"],
                    "price type": item["price"]["price_type"],
                }
            ),
            MappingType.composite,
        ).key["fuelquest id"]

        now = datetime.now(UTC)

        confirmation = DeliveryConfirmation(
            confirmation_type=ReportType.ORIGINAL,
            equipment=DeliveryEquipment(),
            carrier=this_carrier,
            buyer=buyer,
            consignor=Consignor(
                name=buyer_name,  # TODO: Should be store name
                address=store["address"],
                city=store["city"],
                state=store["state"],
                postal_code=store["postal_code"],
                country_code=CountryCode(code=self.country_code),
                organization_id=OrganizationId(
                    ident=buyer_fq_id,
                ),
            ),
            bill_to=BillTo(
                name=buyer_name,
                address=store["address"],
                city=store["city"],
                state=store["state"],
                country_code=CountryCode(code=self.country_code),
                organization_id=OrganizationId(
                    ident=buyer_fq_id,
                ),
            ),
            bill_of_lading_number=item["bol_number"],
            purchase_order_number=po_number or "",
            release_number=ReleaseNumber(value="", expiration_datetime=None),
            delivery_information=DeliveryInformation(
                delivered_to=DeliveredTo(
                    ident=store_fq_id,
                    name=store["store_number"],
                    address=store["address"],
                    city=store["city"],
                    state=store["state"],
                    country_code=CountryCode(
                        code=self.country_code
                    ),
                    organization_id=OrganizationId(
                        ident=store_fq_id,
                    ),
                    # TODO: Always convert TZ to GMT / UTC here. Get real delivery date and times, not "now"
                    start_delivery_datetime=now,
                    end_delivery_datetime=now,
                    carrier_way_bill_number="",
                    line_items=[
                        self._build_delivery_line_item(item, fq_product_id, item_count)
                    ],
                ),
            ),
        )

        invoice_detail = self._build_freight_invoice_detail(
            item=item,
            invoices=invoices,
            order_number=order["order_number"],
            terminal=supply_terminal,
            store=store,
            store_fq_id=store_fq_id,
            supplier_cp_id=item["bol_supplier_id"],
            fq_product_id=fq_product_id,
        )

        return confirmation, invoice_detail

    def _market_by_store_id(self, store_id: str) -> str | None:
        return (self.stores_by_id[store_id] or {}).get("market")

    def _build_delivery_report(
        self, order_list: list[dict], invoices: list[dict]
    ) -> Iterator[DeliveryReport]:
        for order in order_list:
            confirmations = []
            invoice_details: list[FreightInvoiceDetail] = []
            first_carrier = None
            first_buyer = None

            try:
                order = PathDict(order)
                order_number = order["order_number"]

                any_enabled_buyers = any(
                    self._buyer_from_allocated_bol(b) in self.allowed_counterparties
                    for b in order["allocated_bols"] or []
                )
                if not any_enabled_buyers:
                    # logger.info(f"Skipping order {order_number} because it is not for any enabled buyers")
                    self.skipped_orders_filter.append((str(order_number), "No BOLs in order are for enabled buyers"))
                    continue
                logger.info(
                    f"Converting order {order_number} with {len(order['allocated_bols'])} allocated items"
                )

                if self.allowed_store_markets:
                    any_enabled_markets = any(
                        self._market_by_store_id(b["store_id"]) in self.allowed_store_markets
                        for b in order["allocated_bols"] or []
                    )
                    if not any_enabled_markets:
                        self.skipped_orders_filter.append(
                            (str(order_number), "No BOLs in order are for enabled markets"))
                        continue

                for count, item in enumerate(order["allocated_bols"] or []):
                    item = pd(item)
                    po_number = order["order_number"]

                    market = self._market_by_store_id(item["store_id"])
                    if self.allowed_store_markets and market not in self.allowed_store_markets:
                        logger.info(
                            f"Skipping BOL {item['bol_number']} on {order['order_number']}: store market {market} "
                            f"not in allowed store markets ({self.allowed_store_markets})"
                        )
                        continue

                    # Resolve / map the depot from that on the driver schedule
                    depot_grav_name = self.driver_sched_lkp[item["driver_schedule"]][
                        "depot_name"
                    ]
                    depot_fq_id = self.rita_mapper.get_source_parent_id(
                        depot_grav_name, MappingType.depot
                    )
                    truck_terminal = TruckTerminal(
                        name=depot_grav_name,
                        organization_id=OrganizationId(ident=depot_fq_id),
                    )

                    bol_terminal_name = item["bol_terminal"]
                    supply_terminal_fq_id = self.rita_mapper.get_source_parent_id(
                        bol_terminal_name, MappingType.terminal
                    )
                    supply_terminal = SupplyTerminal(
                        name=bol_terminal_name,
                        organization_id=OrganizationId(ident=supply_terminal_fq_id),
                    )

                    result = self._build_confirmation_and_invoice_detail(
                        order,
                        item,
                        invoices,
                        po_number,
                        truck_terminal,
                        supply_terminal,
                        count + 1,
                    )
                    if result is None:
                        continue
                    confirmation, invoice_detail = result
                    confirmations.append(confirmation)
                    invoice_details.append(invoice_detail)
                    if first_carrier is None:
                        first_carrier = confirmation.carrier
                        first_buyer = confirmation.buyer

                if not confirmations:
                    logger.warning(
                        f"No contents converted for order {order_number}, skipping"
                    )
                    continue

                # Find invoice-level fields from the first matching invoice
                invoice_number = "0"
                invoice_date = None
                for inv in invoices:
                    for txn in inv.get("transactions", []):
                        if str(txn["order_number"]) == order_number:
                            invoice_number = inv["invoice_number"]
                            invoice_date = datetime.fromisoformat(
                                inv["created_on"]
                            ).date()
                            break
                    if invoice_date:
                        break

                all_freight_details = [
                    fd for detail in invoice_details for fd in detail.freight_details
                ]

                freight_invoice = FuelFreightInvoice(
                    invoice_type=ReportType.ORIGINAL,
                    carrier=first_carrier,
                    buyer=first_buyer,
                    invoice_number=invoice_number,
                    invoice_date=invoice_date or date.today(),
                    freight_invoice_detail=invoice_details,
                    freight_invoice_summary=FreightInvoiceSummary(
                        totals=FreightInvoiceTotals(
                            freight_details=FreightInvTotalsDetails(
                                count=len(all_freight_details),
                            ),
                            total_due_amount=TotalFreightInvoiceDueAmt(
                                amount=sum(
                                    fd.freight_charge.amount
                                    for fd in all_freight_details
                                )
                            ),
                            additional_allowances_and_charges=AdditionalAllowancesAndCharges(
                                total_additional_net=0.0  # TODO
                            ),
                        )
                    ),
                    terms=Terms(),
                )

                logger.info(f"Converted order {order['order_number']}")
                self.successful_order_numbers.append(order_number)

                # May not have generated any confirmations if they were all for non-enabled counterparties
                yield DeliveryReport(
                    reportType=self.report_type,
                    delivery_confirmation=confirmations,
                    fuel_freight_invoice=freight_invoice,
                )
            except RitaMapperKeyError as e:
                logger.warning(e)
                self.issues.append(self._rmke_to_issue(e, order["order_number"]))
            except Exception as e:
                logger.exception(f"Error processing order {order['order_number']}: {e}")

    async def load_ref_data(self):
        with TimedLog("INFO", "Loading counterparties"):
            ref_cp = (await self.sd_client.all_counterparties()).json()
            self.ref_counterparties = pd(ref_cp)
            self.cp_by_id = pd({cp["id"]: cp for cp in ref_cp})

        with TimedLog("INFO", "Loading products"):
            self.ref_products = pd((await self.sd_client.all_products()).json())

        with TimedLog("INFO", "Loading stores"):
            ref_stores = (await self.sd_client.all_stores(include_tanks=False)).json()
            self.ref_stores = pd(ref_stores)
            self.stores_by_id = pd({store["_id"]: store for store in ref_stores})

        with TimedLog("INFO", "Loading locations"):
            ref_locations = (await self.sd_client.all_locations()).json()
            self.ref_locations = pd(ref_locations)
            self.locations_by_id = pd({loc["id"]: loc for loc in ref_locations})

    async def convert(
        self, order_list: list[dict], invoices: list[dict]
    ) -> AsyncIterator[FuelQuestNAXMLDelivery]:
        # Have to init this here since we don't know the driver schedules until we have the order list
        with TimedLog("INFO", "Loading driver schedules"):
            drivers = await self.sd_client.get_driver_tracking(
                order_numbers=[o["order_number"] for o in order_list]
            )
            drivers.raise_for_status()
            self.driver_sched_lkp = pd(
                {driver["driver_schedule_id"]: driver for driver in drivers.json()}
            )

        # TODO: Ensure that the location = consignor
        header = TransmissionHeader(
            transmission_id=self.transmission_datetime.strftime("%Y%m%d%H%M%S"),
            transmission_datetime=self.transmission_datetime,
            transmission_status=self.report_type,
            transmission_sender=self.transmission_sender,
            transmission_receiver=self.transmission_receiver,
        )
        for report in self._build_delivery_report(order_list, invoices):
            yield FuelQuestNAXMLDelivery(
                transmission_header=header, delivery_report=report
            )
