from typing import Any, Dict, List, Optional, Tuple

import pandas as pd
from django.contrib import admin
from django.contrib.admin.views.decorators import staff_member_required
from django.db.models import Case, Count, F, Q, QuerySet, When
from django.http import HttpRequest, HttpResponse
from django.shortcuts import render
from django.utils import timezone

from promptgame.gameui.models import Account, AuthAttempt, Defense, Flag, User
from promptgame.gameui.restrictions import (
    DIFFICULTY_GROUP_MAP,
    RESTRICTIONS_DICT,
    Difficulty,
)


def make_percent(num: float, denom: float, decimal_places: int = 2) -> Optional[float]:
    return round(num / denom * 100, decimal_places) if denom > 0 else None


def get_active_accounts(hours: int) -> QuerySet[Account]:
    now = timezone.now()
    earlier_time = now - timezone.timedelta(hours=hours)

    # 1. Accounts with attack attempts during the time delta
    attacker_ids = AuthAttempt.objects.filter(timestamp__gte=earlier_time).values_list(
        "attacking_account", flat=True
    )

    # 2. Accounts with active defenses during the time delta
    active_defender_ids = Account.objects.filter(
        active_defense__timestamp__gte=earlier_time
    ).values_list("id", flat=True)

    # 3. Combine results and retrieve distinct Account objects
    active_accounts = Account.objects.filter(
        Q(id__in=attacker_ids) | Q(id__in=active_defender_ids)
    ).distinct()
    return active_accounts


def top_k_active_accounts(k: int, hours: int, model: type) -> List[Account]:
    earlier_time = timezone.now() - timezone.timedelta(hours=hours)

    # Check which model is being used to determine the correct related field for counting
    if model == AuthAttempt:
        related_field = "attack_attempts"
    elif model == Defense:
        related_field = "defenses"
    else:
        raise ValueError("Invalid model passed!")

    # Annotate each Account with the count of related instances from the specified model
    # created after earlier_time
    top_accounts = Account.objects.annotate(
        model_count=Count(
            related_field,
            filter=Q(**{f"{related_field}__timestamp__gte": earlier_time}),
        )
    ).order_by("-model_count")[:k]

    return list(top_accounts)


def random_k_objects(k: int, hours: int, model: Any) -> List[Any]:
    earlier_time = timezone.now() - timezone.timedelta(hours=hours)

    # Filter the model queryset to only include attempts from the last 24 hours
    # and order them randomly, then fetch the first k.
    random_attempts = model.objects.filter(timestamp__gte=earlier_time).order_by("?")[
        :k
    ]

    # Return the selected models
    return list(random_attempts)


def get_time_interval_logs(hours: int) -> Dict[str, Optional[float]]:
    now = timezone.now()
    earlier_time = now - timezone.timedelta(hours=hours)
    defense_count = Defense.objects.filter(timestamp__gte=earlier_time).count()
    auth_attempts = AuthAttempt.objects.filter(timestamp__gte=earlier_time)
    non_self_auth_attempts = auth_attempts.exclude(
        attacking_account=F("defending_account")
    )
    self_auth_attempts = auth_attempts.filter(attacking_account=F("defending_account"))
    successful_non_self_auth_attempts = non_self_auth_attempts.filter(
        was_successful=True
    )
    unique_non_self_auth_attempts = non_self_auth_attempts.values_list(
        "attacker_input", flat=True
    ).distinct()
    unique_successful_non_self_auth_attempts = (
        successful_non_self_auth_attempts.values_list(
            "attacker_input", flat=True
        ).distinct()
    )
    unique_attackers = (
        AuthAttempt.objects.filter(timestamp__gte=earlier_time)
        .values_list("attacking_account", flat=True)
        .distinct()
    )
    unique_successful_attackers = (
        AuthAttempt.objects.filter(timestamp__gte=earlier_time, was_successful=True)
        .values_list("attacking_account", flat=True)
        .distinct()
    )

    auth_attempt_count = auth_attempts.count()
    non_self_auth_attempt_count = non_self_auth_attempts.count()
    successful_non_self_auth_attempt_count = successful_non_self_auth_attempts.count()
    unique_successful_non_self_auth_attempt_count = (
        unique_successful_non_self_auth_attempts.count()
    )
    flag_count = Flag.objects.filter(timestamp__gte=earlier_time).count()
    unique_attackers_count = unique_attackers.count()
    unique_successful_attackers_count = unique_successful_attackers.count()
    unique_non_self_auth_attempt_count = unique_non_self_auth_attempts.count()
    self_attack_count = self_auth_attempts.count()
    new_users_count = User.objects.filter(date_joined__gte=earlier_time).count()

    non_self_non_moderated_authattempts = non_self_auth_attempts.filter(
        timestamp__gte=earlier_time, auto_moderated=False
    ).count()
    non_moderated_defenses = Defense.objects.filter(
        timestamp__gte=earlier_time, auto_moderated=False
    ).count()

    active_accounts_count = get_active_accounts(hours).count()
    authattempts_with_restrictions_count = auth_attempts.exclude(
        restriction_id=None
    ).count()
    successful_authattempts_with_restrictions_count = (
        auth_attempts.exclude(restriction_id=None).filter(was_successful=True).count()
    )

    return {
        "active_accounts": active_accounts_count,
        "active_percent": make_percent(active_accounts_count, Account.objects.count()),
        "users_per_hr": round(new_users_count / hours, 2),
        "flags_per_hr": round(flag_count / hours, 2),
        "defenses_per_hr": round(defense_count / hours, 2),
        "nonself_authattempts_per_hr": round(non_self_auth_attempt_count / hours, 2),
        "successful_nonself_authattempts_per_hr": round(
            successful_non_self_auth_attempt_count / hours, 2
        ),
        "percent_selfattack": make_percent(self_attack_count, auth_attempt_count),
        "percent_nonself_authattempts_successful": make_percent(
            successful_non_self_auth_attempt_count, non_self_auth_attempt_count
        ),
        "percent_nonself_authattempt_unique": make_percent(
            unique_non_self_auth_attempt_count, non_self_auth_attempt_count
        ),
        "percent_successful_nonself_authattempts_unique": make_percent(
            unique_successful_non_self_auth_attempt_count,
            successful_non_self_auth_attempt_count,
        ),
        "unique_attackers": unique_attackers_count,
        "unique_successful_attackers": unique_successful_attackers_count,
        "unmoderated_percent": make_percent(
            non_self_non_moderated_authattempts + non_moderated_defenses,
            non_self_auth_attempt_count + defense_count,
        ),
        "restrictions_percent": make_percent(
            authattempts_with_restrictions_count, auth_attempt_count
        ),
        "restricted_success_percent": make_percent(
            successful_authattempts_with_restrictions_count,
            authattempts_with_restrictions_count,
        ),
    }


def make_restrictions_html() -> Tuple[pd.DataFrame, str]:
    # Count Accounts for each restriction_id
    account_counts = Account.objects.values("restriction_id").annotate(
        count=Count("id")
    )
    total_accounts = Account.objects.count()

    # Count AuthAttempt for each restriction_id and success rates
    auth_counts = (
        AuthAttempt.objects.exclude(restriction_id=None)
        .values("restriction_id")
        .annotate(
            count=Count("id"),
            success_count=Count(Case(When(was_successful=True, then=1))),
        )
    )
    total_auth_attempts_with_restrictions = AuthAttempt.objects.exclude(
        restriction_id=None
    ).count()

    def make_restriction_id_dict(query_set, key="restriction_id", val="count"):
        # Turn a QuerySet into a dictionary {key: val}, where key and val are both
        # attributes of each item in the query_set.
        return {item[key]: item[val] for item in query_set}

    account_counts_dict = make_restriction_id_dict(
        account_counts, "restriction_id", "count"
    )
    auth_counts_dict = make_restriction_id_dict(auth_counts, "restriction_id", "count")
    success_counts_dict = make_restriction_id_dict(
        auth_counts, "restriction_id", "success_count"
    )

    # Find the fraction of restriction_ids in each difficulty bucket
    difficulty_dict = {bucket: 0 for bucket in Difficulty}
    for item in account_counts:
        # Get the difficulty of the restriction_id
        difficulty = RESTRICTIONS_DICT[item["restriction_id"]].difficulty
        # Increment the count for that difficulty
        difficulty_dict[difficulty] += item["count"]

    def get_expected_weight(restriction_group):
        # We sample by taking a uniform sample over each restriction_group in the difficulty group
        # and then uniformly sampling over the restriction_ids in that restriction_group.
        # The expected weight of a restriction_group is a sum(group_frequency * 1/group_size if group in difficulty else 0).

        expected_weight = 0
        for difficulty, difficulty_count in difficulty_dict.items():
            group_fraction = difficulty_count / total_accounts
            if restriction_group in DIFFICULTY_GROUP_MAP[difficulty]:
                expected_weight += group_fraction / len(
                    DIFFICULTY_GROUP_MAP[difficulty]
                )

        return expected_weight

    def get_group(restriction_id):
        if restriction_id in RESTRICTIONS_DICT:
            return RESTRICTIONS_DICT[restriction_id].restriction_group
        # might happen when it's an old restriction that we have since removed
        return "unknown"

    # Initialize an empty DataFrame
    columns = [
        "restriction_group",
        "account %",
        "expected %",
        "representation %",
        "auth_attempt %",
        "success %",
        "groups",
    ]
    df = pd.DataFrame(columns=columns)

    # Make a list of all restriction_ids for each restriction_group.
    all_restriction_groups = set()
    for restriction_groups in DIFFICULTY_GROUP_MAP.values():
        all_restriction_groups.update(restriction_groups)

    n = len(all_restriction_groups) + 1  # +1 for the 'other' category
    df = pd.DataFrame(index=range(n), columns=columns)

    def sum_values_for_restrictions_in_group(d, restriction_group):
        # Takes in a dict where restriction ids are keys.
        # Outputs a the sum of the values from restriction_ids in the group.
        total = 0
        for restriction_id, v in d.items():
            if get_group(restriction_id) == restriction_group:
                total += v
        return total

    # Fill in DataFrame
    for i, restriction_group in enumerate(all_restriction_groups):
        # TODO: the way this is called makes it quadratic in the number of
        # restrictions. If restriction set grows, try making it linear instead.
        account_count = sum_values_for_restrictions_in_group(
            account_counts_dict, restriction_group
        )
        auth_count = sum_values_for_restrictions_in_group(
            auth_counts_dict, restriction_group
        )
        success_count = sum_values_for_restrictions_in_group(
            success_counts_dict, restriction_group
        )
        account_percent = make_percent(account_count, total_accounts)
        auth_attempt_percent = make_percent(
            auth_count, total_auth_attempts_with_restrictions
        )
        success_percent = make_percent(success_count, auth_count)

        # Representation percent is (fraction of accounts with this restriction) / (expected fraction based on weights)
        expected_weight = get_expected_weight(restriction_group)
        real_weight = account_count / total_accounts
        representation_percent = make_percent(real_weight, expected_weight)

        df.loc[i, "restriction_group"] = restriction_group
        df.loc[i, "account %"] = account_percent
        df.loc[i, "expected %"] = make_percent(expected_weight, 1)
        df.loc[i, "representation %"] = representation_percent
        df.loc[i, "auth_attempt %"] = auth_attempt_percent
        df.loc[i, "success %"] = success_percent
        df.loc[i, "groups"] = ", ".join(
            [
                difficulty.name
                for difficulty, restriction_groups in DIFFICULTY_GROUP_MAP.items()
                if restriction_group in restriction_groups
            ]
        )

    # Calculate 'other' values
    account_count_other = sum(
        item["count"]
        for item in account_counts
        if item["restriction_id"] not in RESTRICTIONS_DICT.keys()
    )
    auth_count_other = sum(
        item["count"]
        for item in auth_counts
        if item["restriction_id"] not in RESTRICTIONS_DICT.keys()
    )
    success_count_other = sum(
        item["success_count"]
        for item in auth_counts
        if item["restriction_id"] not in RESTRICTIONS_DICT.keys()
    )

    account_percent_other = make_percent(account_count_other, total_accounts)
    auth_attempt_percent_other = make_percent(
        auth_count_other, total_auth_attempts_with_restrictions
    )
    success_percent_other = make_percent(success_count_other, auth_count_other)

    df.loc[n - 1, "restriction_group"] = "other"
    df.loc[n - 1, "account %"] = account_percent_other
    df.loc[n - 1, "expected %"] = 0
    df.loc[n - 1, "representation %"] = 0
    df.loc[n - 1, "auth_attempt %"] = auth_attempt_percent_other
    df.loc[n - 1, "success %"] = success_percent_other
    df.loc[n - 1, "groups"] = "N/A"

    html_table = df.to_html(
        index=False, escape=True, float_format=lambda x: f"{x:.2f}%"
    )
    return df, html_table


@staff_member_required
def basic_moderation_stats(request: HttpRequest) -> HttpResponse:
    """Basic statistics for moderators to track app usage."""
    data_dict = {}

    # times to log: 1 day, 3 days, 1 week
    for num_hours in [24, 72, 24 * 7]:
        data_dict.update(
            {
                f"{k}_{num_hours}": v
                for k, v in get_time_interval_logs(num_hours).items()
            }
        )

    count = 10
    time_interval = 24  # 1 day

    # Most active users
    top_k_active_attackers = top_k_active_accounts(count, time_interval, AuthAttempt)
    top_k_active_defenders = top_k_active_accounts(count, time_interval, Defense)

    # Random AuthAttempts and Defenses from the last 24 hours
    random_authattempts = random_k_objects(count, time_interval, AuthAttempt)
    random_defenses = random_k_objects(count, time_interval, Defense)

    # Active accounts
    active_accounts = get_active_accounts(time_interval)
    active_unlocked_accounts = active_accounts.filter(unlock_at__lte=timezone.now())

    _, restrictions_html = make_restrictions_html()

    return render(
        request,
        "admin/promptgame/basic_moderation_stats.html",
        {
            # Recent Defenses and AuthAttempts table
            "recent_defenses": random_defenses,
            "recent_auth_attempts": random_authattempts,
            # Most Active tables
            "most_active_accounts_authattempts": top_k_active_attackers,
            "most_active_accounts_defenses": top_k_active_defenders,
            # Totals table
            "total_accounts": Account.objects.count(),
            "total_defenses": Defense.objects.count(),
            "total_auth_attempts": AuthAttempt.objects.count(),
            "total_successful_auth_attempts": AuthAttempt.objects.filter(
                was_successful=True
            ).count(),
            "total_non_self_auth_attempts": AuthAttempt.objects.exclude(
                attacking_account=F("defending_account")
            ).count(),
            "total_successful_non_self_auth_attempts": AuthAttempt.objects.filter(
                was_successful=True
            )
            .exclude(attacking_account=F("defending_account"))
            .count(),
            "total_flags": Flag.objects.count(),
            "total_unreviewed_flags": Flag.objects.filter(
                reviewed_by_admin=False
            ).count(),
            "unlocked_percent": make_percent(
                Account.objects.filter(unlock_at__lte=timezone.now()).count(),
                Account.objects.count(),
            ),
            "active_unlocked_percent": make_percent(
                active_unlocked_accounts.count(), active_accounts.count()
            ),
            "restrictions_html": restrictions_html,
            # Balance and usage stats table
            **data_dict,
            **admin.site.each_context(request),
        },
    )
