from __future__ import annotations

from dataclasses import replace
from decimal import Decimal, ROUND_FLOOR, ROUND_HALF_UP

from .models import (
    Account,
    BillingEntry,
    Cap,
    Credit,
    Invoice,
    LineItem,
    Period,
    Promotion,
    RuleSet,
    Tax,
    UsageLine,
)

CENT = Decimal("0.01")


def money(value: Decimal) -> Decimal:
    return value.quantize(CENT, rounding=ROUND_HALF_UP)


def allocate(
    amount: Decimal, weights: list[Decimal], keys: list[str]
) -> list[Decimal]:
    amount = money(amount)
    if not weights:
        return []
    total_weight = sum(weights, Decimal("0"))
    if amount == 0 or total_weight == 0:
        return [Decimal("0.00") for _ in weights]
    total_cents = int(amount / CENT)
    quotas = [Decimal(total_cents) * weight / total_weight for weight in weights]
    cents = [
        int(quota.to_integral_value(rounding=ROUND_FLOOR)) for quota in quotas
    ]
    remaining = total_cents - sum(cents)
    order = sorted(
        range(len(weights)),
        key=lambda index: (-(quotas[index] - cents[index]), keys[index]),
    )
    for index in order[:remaining]:
        cents[index] += 1
    return [Decimal(value) * CENT for value in cents]


class PricingEngine:
    def _usage_amount(self, usage: UsageLine, rules: RuleSet) -> Decimal:
        remaining = usage.quantity
        lower = Decimal("0")
        amount = Decimal("0")
        for tier in rules.usage_tiers:
            if remaining <= 0:
                break
            units = (
                remaining
                if tier.up_to is None
                else min(remaining, tier.up_to - lower)
            )
            if units > 0:
                amount += units * tier.unit_price
                remaining -= units
            if tier.up_to is not None:
                lower = tier.up_to
        return amount

    def _pipeline(
        self,
        bases: list[Decimal],
        keys: list[str],
        rules: tuple[object, ...],
        selected: Promotion | None,
    ) -> tuple[list[Decimal], list[Decimal], Decimal]:
        promotions = sorted(
            (
                rule
                for rule in rules
                if isinstance(rule, Promotion) and rule.kind == "stackable"
            ),
            key=lambda rule: rule.key,
        )
        if selected is not None:
            promotions.append(selected)
        raw_reductions: list[Decimal] = []
        for key, base in zip(keys, bases):
            factor = Decimal("1")
            for promotion in promotions:
                if (
                    promotion.applies_to_usage_keys is None
                    or key in promotion.applies_to_usage_keys
                ):
                    factor *= Decimal("1") - promotion.percent
            raw_reductions.append(base * (Decimal("1") - factor))
        reductions = allocate(
            money(sum(raw_reductions, Decimal("0"))), raw_reductions, keys
        )
        nets = [base - reduction for base, reduction in zip(bases, reductions)]

        credit = sum(
            (rule.amount for rule in rules if isinstance(rule, Credit)),
            Decimal("0"),
        )
        subtotal = sum(nets, Decimal("0"))
        credits = allocate(min(credit, subtotal), nets, keys)
        nets = [net - part for net, part in zip(nets, credits)]

        caps = [rule.amount for rule in rules if isinstance(rule, Cap)]
        subtotal = sum(nets, Decimal("0"))
        if caps and subtotal > min(caps):
            parts = allocate(subtotal - min(caps), nets, keys)
            nets = [net - part for net, part in zip(nets, parts)]

        rate = sum(
            (rule.rate for rule in rules if isinstance(rule, Tax)), Decimal("0")
        )
        taxes = allocate(money(sum(nets, Decimal("0")) * rate), nets, keys)
        total = money(sum(nets, Decimal("0")) + sum(taxes, Decimal("0")))
        return nets, taxes, total

    def _build(
        self,
        account_id: str,
        period: Period,
        usage: tuple[UsageLine, ...],
        rules: RuleSet,
        invoice_id: str,
        revision_of: str | None = None,
    ) -> Invoice:
        keys = [line.key for line in usage]
        raw_bases = [self._usage_amount(line, rules) for line in usage]
        bases = allocate(money(sum(raw_bases, Decimal("0"))), raw_bases, keys)
        exclusive = [
            rule
            for rule in rules.rules
            if isinstance(rule, Promotion)
            and rule.kind == "exclusive"
            and (
                rule.applies_to_usage_keys is None
                or set(keys).intersection(rule.applies_to_usage_keys)
            )
        ]
        outcomes = [
            (*self._pipeline(bases, keys, rules.rules, item), item)
            for item in exclusive
        ]
        if outcomes:
            nets, taxes, total, selected = min(
                outcomes, key=lambda result: (result[2], result[3].key)
            )
        else:
            nets, taxes, total = self._pipeline(bases, keys, rules.rules, None)
            selected = None
        items = tuple(
            LineItem(
                line.key,
                line.description,
                line.quantity,
                base,
                net,
                tax,
                money(net + tax),
            )
            for line, base, net, tax in zip(usage, bases, nets, taxes)
        )
        return Invoice(
            invoice_id,
            account_id,
            period,
            usage,
            rules,
            items,
            total,
            selected.key if selected else None,
            tuple(sorted(item.key for item in exclusive if item is not selected)),
            revision_of=revision_of,
        )

    def price_period(
        self,
        account: Account,
        period: Period,
        usage: tuple[UsageLine, ...],
        rule_set: RuleSet,
    ) -> Invoice:
        invoice = self._build(
            account.account_id,
            period,
            usage,
            rule_set,
            f"INV-{account.account_id}-{period.key}",
        )
        account.invoices[period.key] = invoice
        account._base_states[period.key] = (period, usage, rule_set)
        return invoice

    def _active_adjustments(
        self, account: Account, cutoff: int | None
    ) -> list[object]:
        accepted = [
            item
            for item in account.received_adjustments
            if cutoff is None
            or account.acceptance_ordinals[item.adjustment_id] <= cutoff
        ]
        withdrawals = {
            item.withdraws_adjustment_id
            for item in accepted
            if item.kind == "withdrawal"
        }
        nodes = {
            item.adjustment_id: item
            for item in accepted
            if item.kind != "withdrawal"
        }
        children = {
            item.replaces_adjustment_id: item.adjustment_id
            for item in nodes.values()
            if item.replaces_adjustment_id is not None
        }
        roots = [
            item
            for item in nodes.values()
            if item.replaces_adjustment_id is None
        ]
        active = []
        for root in roots:
            chain = [root]
            while chain[-1].adjustment_id in children:
                chain.append(nodes[children[chain[-1].adjustment_id]])
            chosen = next(
                (
                    item
                    for item in reversed(chain)
                    if item.adjustment_id not in withdrawals
                ),
                None,
            )
            if chosen is not None:
                active.append(chosen)
        return sorted(active, key=lambda item: (item.effective_at, item.adjustment_id))

    def resolved_invoice(
        self, account: Account, period_key: str, cutoff: int | None = None
    ) -> Invoice:
        period, base_usage, base_rules = account._base_states[period_key]
        usage = list(base_usage)
        rules = list(base_rules.rules)
        for adjustment in self._active_adjustments(account, cutoff):
            affects = adjustment.period_key == period_key or (
                adjustment.kind == "usage_transfer"
                and adjustment.target_period_key == period_key
            )
            if not affects:
                continue
            if adjustment.kind == "usage_correction":
                usage = [
                    replace(line, quantity=adjustment.quantity)
                    if line.key == adjustment.usage_key
                    else line
                    for line in usage
                ]
            elif adjustment.kind == "usage_transfer":
                moved = adjustment.quantity or Decimal("0")
                sign = -moved if adjustment.period_key == period_key else moved
                usage = [
                    replace(line, quantity=line.quantity + sign)
                    if line.key == adjustment.usage_key
                    else line
                    for line in usage
                ]
            elif adjustment.kind == "credit":
                rules.append(
                    Credit(
                        f"adjustment:{adjustment.adjustment_id}",
                        adjustment.amount or Decimal("0.00"),
                    )
                )
            elif (
                adjustment.kind == "retroactive_promotion"
                and adjustment.promotion is not None
            ):
                rules.append(adjustment.promotion)
        resolved_rules = replace(base_rules, rules=tuple(rules))
        original = account.invoices[period_key]
        return self._build(
            account.account_id,
            period,
            tuple(usage),
            resolved_rules,
            original.invoice_id,
        )

    def issue(self, account: Account, invoice: Invoice) -> None:
        if invoice.closed:
            return
        invoice.closed = True
        invoice.issued_total = invoice.total
        account.billing_entries.append(
            BillingEntry(
                f"BE-{invoice.invoice_id}",
                "invoice",
                invoice.invoice_id,
                invoice.total,
                invoice.period.key,
                invoice.invoice_id,
            )
        )
        account.balance = money(
            sum((entry.amount for entry in account.billing_entries), Decimal("0"))
        )

    def reprice(
        self,
        invoice: Invoice,
        usage: tuple[UsageLine, ...],
        rule_set: RuleSet,
    ) -> Invoice:
        return self._build(
            invoice.account_id,
            invoice.period,
            usage,
            rule_set,
            invoice.invoice_id,
        )