from datetime import date, timedelta
from datetime import date as date_lib
from decimal import Decimal

from django.db import models
from django.db.models import (
    DateField,
    Exists,
    ExpressionWrapper,
    OuterRef,
    Q,
    QuerySet,
    Sum,
)
from django.db.models.functions import Greatest
from django.dispatch import receiver
from django.utils.translation import gettext_lazy as _
from django_fsm import FSMField, transition
from wbcore.contrib.ai.llm.config import add_llm_prompt
from wbcore.contrib.authentication.models import User
from wbcore.contrib.currency.models import CurrencyFXRates
from wbcore.contrib.directory.models import Entry
from wbcore.contrib.icons import WBIcon
from wbcore.enums import RequestType
from wbcore.metadata.configs.buttons import ActionButton, ButtonDefaultColor
from wbcore.metadata.configs.display.instance_display.shortcuts import (
    create_simple_display,
)
from wbcore.models import WBModel
from wbcore.signals import pre_merge
from wbcore.signals.models import pre_collection
from wbcore.utils.enum import ChoiceEnum
from wbcore.utils.strings import ReferenceIDMixin
from wbcrm.models import AccountRole
from wbcrm.models.accounts import Account
from wbfdm.models.instruments.instrument_prices import InstrumentPrice

from ..custodians import Custodian
from ..llm.wbcrm.analyze_relationship import get_performances_prompt
from ..roles import PortfolioRole
from .trades import Trade


def can_administrate_claim(claim, user):
    today = date_lib.today()
    return (
        user.profile.portfolio_roles.filter(
            Q(role_type=PortfolioRole.RoleType.MANAGER)
            & (Q(start__isnull=True) | Q(start__lte=today))
            & (Q(end__isnull=True) | Q(end__gte=today))
        ).exists()
        or user.is_superuser
    )


def can_edit_claim(claim, user):
    return (claim.claimant and claim.claimant.id == user.profile.id) or user.is_internal or user.is_superuser


class ClaimGroupbyChoice(ChoiceEnum):
    ROOT_ACCOUNT = "Root Account"
    ACCOUNT = "Account"
    ROOT_ACCOUNT_OWNER = "Root Account Owner"
    ACCOUNT_OWNER = "Account Owner"
    PRODUCT = "Product"
    PRODUCT_GROUP = "ProductGroup"
    CLASSIFICATION = "Classification"

    @classmethod
    @property
    def map(cls):
        return {
            "ROOT_ACCOUNT": {
                "pk": "root_account",
                "title_key": "root_account_repr",
                "search_fields": ["root_account_repr"],
            },
            "ACCOUNT": {
                "pk": "account",
                "title_key": "account__computed_str",
                "search_fields": ["account__computed_str"],
            },
            "ROOT_ACCOUNT_OWNER": {
                "pk": "root_account_owner",
                "title_key": "root_account_owner_repr",
                "search_fields": ["root_account_owner_repr"],
            },
            "ACCOUNT_OWNER": {
                "pk": "account__owner",
                "title_key": "account__owner__computed_str",
                "search_fields": ["account__owner__computed_str"],
            },
            "PRODUCT": {
                "pk": "product",
                "title_key": "product__computed_str",
                "search_fields": ["product__computed_str"],
            },
            "PRODUCT_GROUP": {
                "pk": "product__parent",
                "title_key": "product__parent__name",
                "search_fields": ["product__parent__name"],
            },
            "CLASSIFICATION": {
                "pk": "classification_id",
                "title_key": "classification_title",
                "search_fields": ["classification_title"],
            },
        }

    @classmethod
    def get_map(cls, name):
        return cls.map[name]


class ClaimDefaultQueryset(QuerySet):
    def filter_for_user(self, user: User, validity_date: date_lib | None = None) -> QuerySet:
        """
        Protect the chained queryset and filter the claims that this user cannot see based on the following rules:
        """
        if user.has_perm("wbcrm.administrate_account"):
            return self
        allowed_accounts = Account.objects.filter_for_user(user, validity_date=validity_date).values(
            "id"
        )  # This is way faster
        # accounts = self.annotate(can_see_account=Exists(allowed_accounts.filter(id=OuterRef("account"))))
        if user.is_internal:
            return self.filter(Q(account__isnull=True) | Q(account__in=allowed_accounts))
        return self.filter(Q(account__in=allowed_accounts) | (Q(account__isnull=True) & Q(creator_id=user.profile.id)))

    def filter_for_customer(self, customer: Entry, include_related_roles: bool = False) -> QuerySet:
        """
        Filter the chained queryset to return only the claim that belongs to a certain customer
        """
        customer_account_ids = Account.get_accounts_for_customer(customer).values("id")

        if not include_related_roles:
            return self.filter(account__in=customer_account_ids)
        return self.annotate(
            role_exists=Exists(AccountRole.objects.filter(entry=customer, account=OuterRef("account")))
        ).filter(Q(account__in=customer_account_ids) | Q(role_exists=True))

    def annotate_asset_under_management_for_date(self, val_date: date):
        return self.annotate(
            net_value=InstrumentPrice.subquery_closest_value("net_value", val_date, instrument_pk_name="product"),
            fx_rate=CurrencyFXRates.get_fx_rates_subquery(val_date),
            asset_under_management=models.ExpressionWrapper(
                models.F("net_value") * models.F("shares"),
                output_field=models.DecimalField(max_digits=16, decimal_places=4, default=0.0),
            ),
            asset_under_management_usd=models.ExpressionWrapper(
                models.F("asset_under_management") * models.F("fx_rate"),
                output_field=models.DecimalField(max_digits=16, decimal_places=4, default=0.0),
            ),
        )


class ClaimManager(models.Manager):
    def get_queryset(self) -> ClaimDefaultQueryset:
        return ClaimDefaultQueryset(self.model)

    def filter_for_user(self, user: User, validity_date: date_lib | None = None) -> QuerySet:
        return self.get_queryset().filter_for_user(user, validity_date=validity_date)

    def filter_for_customer(self, customer: Entry, include_related_roles: bool = False) -> QuerySet:
        return self.get_queryset().filter_for_customer(customer)

    def annotate_asset_under_management_for_date(self, val_date: date) -> QuerySet:
        return self.get_queryset().annotate_asset_under_management_for_date(val_date)


# @workflow.register(serializer_class="wbportfolio.serializers.transactions.claim.ClaimModelSerializer") #we don't register for now. Uncomment as soon as we want to enable workflow on that model
class Claim(ReferenceIDMixin, WBModel):
    """A customer can claim that a trade or part of a trade was executed my them"""

    class Status(models.TextChoices):
        PENDING = "PENDING", "Pending"
        APPROVED = "APPROVED", "Approved"
        WITHDRAWN = "WITHDRAWN", "Withdrawn"
        DRAFT = "DRAFT", "Draft"
        AUTO_MATCHED = "AUTO_MATCHED", "Auto-Matched"

    status = FSMField(default=Status.DRAFT, choices=Status.choices, verbose_name="Status")
    account = models.ForeignKey(
        to="wbcrm.Account",
        related_name="claims",
        null=True,
        blank=True,
        limit_choices_to=models.Q(is_terminal_account=True),
        on_delete=models.SET_NULL,
        verbose_name="Account",
        # help_text="The account the claim is assigned to. If no sub-account is provided it will be tried to be assigned to a sub account based on the given claimant."
    )

    trade = models.ForeignKey(
        to="wbportfolio.Trade",
        related_name="claims",
        blank=True,
        null=True,
        on_delete=models.PROTECT,  # We protect the claim in case of trade deletion. This needs to be handled upstream
        verbose_name="Trade",
        help_text="Please select a Product first. The customer-trade the claim is consolidated against. The customer-trade and the claim don't necessarily have to have the same date, number of shares, etc.",
    )
    product = models.ForeignKey(
        to="wbportfolio.Product", related_name="claims", on_delete=models.PROTECT, verbose_name="Product"
    )
    claimant = models.ForeignKey(
        to="directory.Entry",
        related_name="claimed",
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        verbose_name="Claimant",
        # help_text="The person / company that claims this trade. If no claimant is given, it will assigned to the current user."
    )
    creator = models.ForeignKey(
        to="directory.Entry",
        related_name="claims_created",
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        verbose_name="Creator",
    )
    date = models.DateField(verbose_name="Trade Date")
    bank = models.CharField(max_length=255, blank=True, verbose_name="Bank")
    reference = models.CharField(max_length=255, null=True, blank=True, verbose_name="Additional Reference")
    shares = models.DecimalField(
        max_digits=15,
        decimal_places=4,
        help_text="The amount of shares purchased / sold.",
        null=True,
        blank=True,
        verbose_name="Shares",
    )

    nominal_amount = models.DecimalField(
        max_digits=15,
        decimal_places=2,
        help_text="The nominal amount purchased / sold. Either shares or nominal_amount has to be provided.",
        null=True,
        blank=True,
        verbose_name="Nominal Amount",
    )
    external_id = models.CharField(
        max_length=255,
        null=True,
        blank=True,
        help_text="An External ID to reference a claim.",
        verbose_name="External Identifier",
    )
    as_shares = models.BooleanField(default=True)

    objects = ClaimManager()

    class Meta:
        verbose_name = "Claim"
        verbose_name_plural = "Claims"

    def __str__(self):
        return f"{self.reference_id} {self.product.name} ({self.bank} - {self.shares:,} shares - {self.date}) "

    def save(self, *args, auto_match: bool = True, **kwargs):
        if self.shares is None and self.nominal_amount is None:
            raise ValueError(
                f"Either shares or nominal amount have to be provided. Shares={self.shares}, Nominal={self.nominal_amount}"
            )
        if self.product:
            if self.shares is not None:
                self.nominal_amount = self.shares * self.product.share_price
            else:
                self.shares = self.nominal_amount / self.product.share_price
        if not self.trade and self.status == self.Status.DRAFT and auto_match:
            self.trade = self.auto_match()
        # we shouldn't remove the link to the trade even if the claim is withdrawn for audit purposes, this might screw any reverse queryset such as claims.all() if widthdrawn claims are not explicitly excluded.
        # if self.status == self.Status.WITHDRAWN and self.trade:
        #     self.trade = None
        if self.trade:
            if not self.trade.is_claimable:
                raise ValueError("The selected trade is not a valid customer trade")
            if self.product and self.trade.underlying_instrument.id != self.product.id:
                raise ValueError("The selected product does not match the trade underlying instrument")
        super().save(*args, **kwargs)

    @classmethod
    def get_valid_and_approved_claims(cls, val_date: date_lib | None = None, account: Account | None = None):
        claims = cls.objects.annotate(
            date_considered=ExpressionWrapper(
                Greatest("trade__transaction_date", "date") + 1, output_field=DateField()
            )
        ).filter(Q(account__is_terminal_account=True) & Q(account__is_active=True) & Q(status=cls.Status.APPROVED))
        if val_date:
            claims = claims.filter(date_considered__lt=val_date)
        if account:
            claims = claims.filter(account__in=account.get_descendants(include_self=True))
        return claims

    @property
    def nominal_value(self):
        """Returns the nominal value of a claim

        Shares x Share Price

        Returns:
            Decimal -- Nominal Value
        """
        return self.shares * self.product.share_price

    @transition(
        status,
        [Status.PENDING, Status.AUTO_MATCHED],
        Status.APPROVED,
        permission=can_administrate_claim,
        custom={
            "_transition_button": ActionButton(
                method=RequestType.PATCH,
                identifiers=("wbportfolio:claim",),
                icon=WBIcon.APPROVE.icon,
                color=ButtonDefaultColor.SUCCESS,
                key="approve",
                label="Approve",
                action_label="Approve",
                description_fields="<p>Date: {{date}}</p><p>Shares: {{shares}}</p><p>Product: {{_product.name}} ({{_product.isin}})</p>",
                instance_display=create_simple_display(
                    [
                        ["date", "shares", "bank"],
                        ["product", "claimant", "."],
                        ["account", "account", "account"],
                        ["trade", "trade", "trade"],
                    ]
                ),
            )
        },
    )
    def approve(self, **kwargs):
        pass

    def can_approve(self):
        errors = dict()

        if not self.trade:
            errors["trade"] = [_("With this status, this has to be provided.")]

        if not self.product:
            errors["product"] = [_("With this status, this has to be provided.")]

        if not self.account:
            errors["account"] = [_("With this status, this has to be provided.")]

        # check if the specified product have a valid nav at the specified date
        if (
            (product := self.product)
            and (claim_date := self.date)
            and not product.valuations.filter(date=claim_date).exists()
        ):
            if (prices_qs := product.valuations.filter(date__lt=claim_date)).exists():
                errors["date"] = [
                    f"For product {product.name}, the latest valid valuation date before {claim_date:%Y-%m-%d} is {prices_qs.latest('date').date:%Y-%m-%d}: Please select a valid date."
                ]
            else:
                errors["date"] = [
                    f"There is no valuation before {claim_date:%Y-%m-%d} for product {product.name}: Please select a valid date."
                ]
        return errors

    @transition(
        status,
        [Status.PENDING, Status.AUTO_MATCHED],
        Status.DRAFT,
        permission=lambda claim, user: user.is_internal or user.is_superuser,
        custom={
            "_transition_button": ActionButton(
                method=RequestType.PATCH,
                color=ButtonDefaultColor.ERROR,
                identifiers=("wbportfolio:claim",),
                icon=WBIcon.UNDO.icon,
                key="backtodraft",
                label="Back to Draft",
                action_label="Back to Draft",
                description_fields="<p>Date: {{date}}</p><p>Shares: {{shares}}</p>",
            )
        },
    )
    def backtodraft(self, **kwargs):
        pass

    @transition(
        status,
        [Status.DRAFT],
        Status.WITHDRAWN,
        permission=can_edit_claim,
        custom={
            "_transition_button": ActionButton(
                method=RequestType.PATCH,
                identifiers=("wbportfolio:claim",),
                color=ButtonDefaultColor.ERROR,
                icon=WBIcon.DELETE.icon,
                key="withdraw",
                label="Withdraw",
                action_label="Withdraw",
                description_fields="<p>Date: {{date}}</p><p>Shares: {{shares}}</p>",
            )
        },
    )
    def withdraw(self, **kwargs):
        pass

    @transition(
        status,
        [Status.APPROVED],
        Status.DRAFT,
        permission=lambda claim, user: user.is_internal or user.is_superuser,
        custom={
            "_transition_button": ActionButton(
                method=RequestType.PATCH,
                identifiers=("wbportfolio:claim",),
                icon=WBIcon.EDIT.icon,
                color=ButtonDefaultColor.WARNING,
                key="revise",
                label="Revise",
                action_label="Revise",
                description_fields="<p>Date: {{date}}</p><p>Product: {{_product.name}}</p>",
            )
        },
    )
    def revise(self, **kwargs):
        pass

    @transition(
        status,
        [Status.DRAFT, Status.AUTO_MATCHED],
        Status.PENDING,
        permission=lambda claim, user: user.is_internal or user.is_superuser,
        custom={
            "_transition_button": ActionButton(
                method=RequestType.PATCH,
                identifiers=("wbportfolio:claim",),
                icon=WBIcon.SEND.icon,
                color=ButtonDefaultColor.WARNING,
                key="submit",
                label="Submit for approval",
                action_label="Submit",
                description_fields="<p>Date: {{date}}</p><p>Product: {{_product.name}}</p>",
            )
        },
    )
    def submit(self, **kwargs):
        pass

    def can_submit(self):
        return self.can_approve()

    def auto_match(self) -> Trade | None:
        shares_epsilon = 1  # share
        auto_match_trade = None
        # Obvious filtering
        trades = Trade.valid_customer_trade_objects.filter(
            Q(transaction_date__gte=self.date - timedelta(days=Trade.TRADE_WINDOW_INTERVAL))
            & Q(transaction_date__lte=self.date + timedelta(days=Trade.TRADE_WINDOW_INTERVAL))
        )
        if self.product:
            trades = trades.filter(underlying_instrument=self.product)
        # Find trades by shares (or remaining to be claimed)
        trades = trades.filter(
            Q(diff_shares__lte=self.shares + shares_epsilon) & Q(diff_shares__gte=self.shares - shares_epsilon)
        )
        if trades.count() == 1:
            auto_match_trade = trades.first()

        # Find trades by custodian
        if (
            not auto_match_trade
            and self.bank
            and (custodian := Custodian.get_by_mapping(self.bank, use_similarity=True, create_missing=False))
        ):
            if trades.filter(custodian=custodian).exists():
                trades = trades.filter(custodian=custodian)
            if trades.count() == 1:
                auto_match_trade = trades.first()

        # Find trades by external_id
        if not auto_match_trade and self.external_id and trades.count() > 1:
            trades = trades.filter(
                Q(external_id__icontains=self.external_id) | Q(external_id_alternative__icontains=self.external_id)
            )
            if trades.count() == 1:
                auto_match_trade = trades.first()

        if auto_match_trade:
            self.status = self.Status.AUTO_MATCHED
            return auto_match_trade

    @classmethod
    def subquery_assets_under_management_for_account(cls, claims, price_date, account_key="account__pk"):
        """Returns a subquery which annotates the assets in USD for each Sub Account based on a given queryset

        Arguments:
            claims {QuerySet<commission.Claim>} -- Prefiltered queryset of claims
            price_date {datetime.date} -- the date which is used to calculate everything
            account_key {str} -- the outer reference to the sub account pk (default: account__pk)

        Returns:
            Subquery -- Sub Account with assets in USD
        """

        return models.Subquery(
            claims.filter(
                status=cls.Status.APPROVED,
                account__pk=models.OuterRef(account_key),
                trade__transaction_date__lte=price_date,
            )
            .annotate_asset_under_management_for_date(price_date)
            .values("account")
            .annotate(sum_assets_usd=models.Sum("assets_under_management_usd"))
            .values("sum_assets_usd")[:1],
            output_field=models.DecimalField(max_digits=16, decimal_places=4, default=0.0),
        )

    @classmethod
    def subquery_claim_sum_for_account(cls, claims, price_date, account_key="account__pk"):
        return models.Subquery(
            claims.filter(
                status=cls.Status.APPROVED,
                account__pk=models.OuterRef(account_key),
                trade__transaction_date__lte=price_date,
            )
            .values("account")
            .annotate(claim_sum=models.Sum("shares"))
            .values("claim_sum")[:1],
            output_field=models.DecimalField(max_digits=16, decimal_places=0, default=0.0),
        )

    @classmethod
    def get_endpoint_basename(cls):
        return "wbportfolio:claim"

    @classmethod
    def get_representation_endpoint(cls):
        return "wbportfolio:claimrepresentation-list"

    @classmethod
    def get_representation_value_key(cls):
        return "id"

    @classmethod
    def get_representation_label_key(cls):
        return "{{shares}} ({{account}})"


@receiver(models.signals.post_save, sender="wbportfolio.Trade")
def post_save_trade(sender, instance, created, raw, **kwargs):
    """
    For every claimable trade, first try to autoclaim the already existing draft claims with this new trade. And auto claim them as well if the product has a default sub account
    """
    if not raw and instance.is_claimable and created:
        min_transaction_date = instance.transaction_date - timedelta(days=Trade.TRADE_WINDOW_INTERVAL)
        max_transaction_date = instance.transaction_date + timedelta(days=Trade.TRADE_WINDOW_INTERVAL)
        unlinked_claims = Claim.objects.filter(
            status=Claim.Status.DRAFT,
            date__gte=min_transaction_date,
            date__lte=max_transaction_date,
            trade__isnull=True,
            product__id=instance.underlying_instrument.id,
        )
        for claim in unlinked_claims.all():
            claim.trade = claim.auto_match()
            claim.save()
        if instance.product and (account := instance.product.default_account):
            shares = instance.shares or Decimal(0)
            claimed_shares = instance.claims.filter(status=Claim.Status.APPROVED).aggregate(s=models.Sum("shares"))[
                "s"
            ] or Decimal(0)
            rest_shares = shares - claimed_shares
            if rest_shares != 0:
                Claim.objects.create(
                    trade=instance,
                    shares=rest_shares,
                    bank=instance.bank,
                    date=instance.transaction_date,
                    product=instance.product,
                    account=account,
                    claimant=account.owner,
                )
            for claim in instance.claims.exclude(Q(status=Claim.Status.APPROVED) | Q(status=Claim.Status.WITHDRAWN)):
                if claim.status == Claim.Status.DRAFT:
                    claim.submit()
                if claim.status == Claim.Status.PENDING:
                    claim.approve()
                claim.save()


@receiver(pre_collection, sender="wbportfolio.Trade")
def pre_trade_deletion(sender, instance, **kwargs):
    # We can't use pre_delete signal here because the collector is collected before the signal and thus, it will
    for claim in instance.claims.filter(
        status__in=[Claim.Status.DRAFT, Claim.Status.AUTO_MATCHED, Claim.Status.WITHDRAWN]
    ).all():
        claim.trade = None
        claim.save(auto_match=False)

    if (instance.marked_for_deletion or instance.pending) and instance.claims.exists():
        # If a default account exists on the trade's product, it means that the product's trade are automatically claims. In that case, we are sure that a valid trade still exists and the marked for deletion's claim can be safely remove
        if instance.product and instance.product.default_account:
            instance.claims.update(status=Claim.Status.WITHDRAWN)
        else:
            similar_trades = Trade.objects.filter(
                underlying_instrument=instance.underlying_instrument,
                portfolio=instance.portfolio,
                marked_for_deletion=True,
                transaction_date=instance.transaction_date,
            )
            # We check if the sum of marked for deletion trades share sums to 0, in that case, we delete them without regards and set their potential claims to draft
            if similar_trades.exists() and similar_trades.aggregate(s=Sum("shares"))["s"] == 0:
                for t in similar_trades.all():
                    t.claims.exclude(status=Claim.Status.WITHDRAWN).update(trade=None, status=Claim.Status.DRAFT)
            elif (other_unclaims_similar_trades := instance.get_alternative_valid_trades()).exists():
                if other_unclaims_similar_trades.count() > 1:
                    other_unclaims_similar_trades = other_unclaims_similar_trades.filter(bank=instance.bank)
                if other_trade := other_unclaims_similar_trades.first():
                    # If we find an unclaim trade with similar attributes, we forward the marked_for_deletion attribute to it, which will be handled/deleted in a next delete iteration
                    instance.claims.exclude(status=Claim.Status.WITHDRAWN).update(trade=other_trade)


@receiver(pre_merge, sender="wbcrm.Account")
def handle_pre_merge_account_for_claim(
    sender: models.Model, merged_object: "Account", main_object: "Account", **kwargs
):
    """
    Simply reassign the claim linked to the merged account to the main account
    """
    Claim.objects.filter(account=merged_object).update(account=main_object, reference=merged_object.reference_id)


@receiver(add_llm_prompt, sender="wbcrm.Account")
def add_holdings_to_account_heat(sender, instance, key, **kwargs):
    if key == "analyze_relationship":
        return get_performances_prompt(instance)
    return []


@receiver(pre_merge, sender="directory.Entry")
def pre_merge_entry_claim(sender: models.Model, merged_object, main_object, **kwargs):
    Claim.objects.filter(claimant=merged_object).update(claimant=main_object)
    Claim.objects.filter(creator=merged_object).update(creator=main_object)
