import asyncio
import enum
import hashlib
import logging
import math
from dataclasses import dataclass
from datetime import timedelta
from functools import lru_cache
from typing import Any, AsyncIterable, Dict, Optional, Tuple

from asgiref.sync import sync_to_async
from django.conf import settings
from django.db import transaction
from django.template import Context, Template, engines
from django.template.loader import render_to_string
from django.utils import timezone
from openai import OpenAIError

from promptgame.gameui import restrictions
from promptgame.gameui.cooldown import accounts_with_cooldown_info
from promptgame.gameui.exceptions import RateLimitException
from promptgame.gameui.llm_wrappers import (
    LLMChoice,
    LLMInput,
    ModerationResponse,
    ModerationResponseInappropriate,
    call_llm_api,
)
from promptgame.gameui.models import Account, AuthAttempt, Flag
from promptgame.gameui.rate_limiting import RateLimitResult, rate_limit_check
from promptgame.gameui.tiers import valid_tier
from promptgame.gameui.user_account_tools import get_unlock_time
from promptgame.gameui.utils import get_number_tokens, sse_message

log = logging.getLogger(__name__)

# we get this for weird prompts like 300x artisanlibs in a row
INSCRUTABLE_MYSTERY_OPENAI_ERROR = (
    "Detected an error in the prompt. Please try again with a different prompt."
)


class DenialReason(enum.Flag):
    """Reason for denying permission to attack."""

    NO_REASON = 0  # This is just here to support bitwise ops.
    ATTACKER_LOCKED = 1
    DEFENDER_LOCKED = 2
    WRONG_TIERS = 4
    COOLDOWN = 8


# string enum containing emoji
class StateEmoji(str, enum.Enum):
    """Emoji to display for account state."""

    LOCKED = "🔒"
    PROTECTED = "🛡️"
    COOLDOWN = "⏳"
    UNLOCKED = "🟢"
    SELF_ATTACK = "👋"

    def __str__(self):
        return self.value


@lru_cache(maxsize=128)
def _cached_string_template(template_source: str) -> Template:
    """Cache the result of Template(template_source)."""
    engine = engines["django"]
    template = engine.from_string(template_source)
    return Template(template_source)


class SSEMessages:
    """For constructing SSE messages that we can show to the user."""

    @staticmethod
    def template(
        template_name: str,
        context: Optional[Dict[str, Any]] = None,
        *,
        swap_spec: str = "beforeend",
    ) -> bytes:
        rendered = render_to_string(template_name, context)
        return sse_message(data=rendered, event=f"hx-swap:{swap_spec}")

    @staticmethod
    def string_template(
        template_source: str,
        context: Optional[Dict[str, Any]] = None,
        *,
        swap_spec: str = "beforeend",
    ) -> bytes:
        # render the raw template source with the given context
        template = _cached_string_template(template_source)
        rich_context = Context(context)
        rendered = template.render(rich_context)
        return sse_message(
            data=rendered,
            event=f"hx-swap:{swap_spec}",
        )

    @staticmethod
    def clear_response_area() -> bytes:
        """Clear the response area by replacing it with nothing."""
        return SSEMessages.string_template("", swap_spec="innerHTML")

    @staticmethod
    def start_llm_response() -> bytes:
        """Start the LLM response by inserting a response bubble."""
        return SSEMessages.template(
            "chat/llm_response_bubble.html",
            context={
                "llm_response": "",
                "response_elt_id": "llm-response-bubble",
            },
            swap_spec="beforeend",
        )

    @staticmethod
    def append_to_llm_response_bubble(delta: str) -> bytes:
        """Append to the previously-inserted response bubble."""
        return SSEMessages.template(
            "chat/llm_delta_fragment.html",
            context={
                "delta": delta,
                "response_elt_id": "llm-response-bubble",
            },
            # OOB swaps still get processed with `hx-swap: none`
            swap_spec="none",
        )

    @staticmethod
    def add_message_bubble(*, type: str, message: str) -> bytes:
        """Show a message to the user."""
        assert type in ("info", "error", "success", "warning"), f"{type=} is invalid"
        return SSEMessages.template(
            template_name="chat/info_message.html",
            context={
                "type": type,
                "message": message,
            },
            swap_spec="beforeend",
        )

    @staticmethod
    def add_tutorial_success_bubble(
        *, step: int, bounty_message: str, is_new_user: bool
    ) -> bytes:
        """Show a message to the user."""
        return SSEMessages.template(
            template_name="chat/tutorial_success_message.html",
            context={
                "next_step": step + 1,
                "total_steps": len(settings.TUTORIAL_ACCOUNTS),
                "bounty_message": bounty_message,
                "is_new_user": is_new_user,
            },
            swap_spec="beforeend",
        )

    @staticmethod
    def trigger_events(*event_names: str) -> bytes:
        for event_name in event_names:
            assert (
                "," not in event_name and "\n" not in event_name
            ), f"{event_name=} contains a comma or newline"
        events = ",".join(event_names)
        return sse_message(event=f"hx-trigger:{events}", data="")


class PermissionGranted:
    """Returned by attack_permission when it's okay to attack."""


@dataclass(frozen=True)
class PermissionDenied:
    """Returned by attack_permission when it's not okay to attack."""

    reason: DenialReason
    user_message: str
    log_message: str


PermissionType = PermissionGranted | PermissionDenied


# Used to represent missing values
_sentinel = object()


def user_attack_permission(*, attacker: Account, defender: Account) -> PermissionType:
    """Check whether attacker has permission to launch an attack on defender.

    This just handles the account-level checks, like looking at locking. It
    doesn't handle attack validation."""

    # Logic:
    #
    # 1. If self-attack: grant permission, this is always okay.
    # 2. Otherwise:
    #    a. If the attacker is locked, we cannot attack.
    #    b. If defender is locked, we cannot attack.
    #    c. If the attacker and defender are in the wrong tiers, we cannot attack.
    #    d. If the attacker is in cooldown w.r.t. the defender, we cannot attack.
    # 3. If we get this far: grant permission to attack!

    if attacker.account_id == defender.account_id:
        # Always permit self-attacks
        return PermissionGranted()

    attacker_in_cooldown = getattr(defender, "attacker_in_cooldown", _sentinel)
    if attacker_in_cooldown is _sentinel:
        raise ValueError(
            f"{defender=} was not annotated with attacker_in_cooldown attribute. Did "
            "you apply annotate_cooldown_info() to the queryset which produced this "
            "defender name?"
        )

    # Otherwise, record all reasons for denial
    denial_reasons = DenialReason.NO_REASON
    # Note that denial_reason and denial_text_emoji get overwritten, so the code
    # below declares an implicit precedence over error messages.
    denial_reason = ""
    denial_text_emoji = ""
    if defender.locked:
        denial_reasons |= DenialReason.DEFENDER_LOCKED
        denial_reason = "it is locked"
        denial_text_emoji = StateEmoji.LOCKED
    if attacker_in_cooldown:
        denial_reasons |= DenialReason.COOLDOWN
        denial_text_emoji = "⏳"
        denial_reason = "you broke into it recently"
    if not valid_tier(attacker.balance, defender.balance):
        denial_reasons |= DenialReason.WRONG_TIERS
        denial_reason = "it is protected (more than one tier below your tier)"
        denial_text_emoji = "🛡️"
    if attacker.locked:
        denial_reasons |= DenialReason.ATTACKER_LOCKED
        denial_text_emoji = StateEmoji.LOCKED
        denial_reason = "your account is locked"

    # No reason for denial, grant permission
    if denial_reasons == DenialReason.NO_REASON:
        return PermissionGranted()

    # Format thse, format the reasons for denial
    denial_text = (
        f"{denial_text_emoji} You can't attack this account because {denial_reason}"
    )
    log_text = f"{attacker.account_id=} denied permission to attack {defender.account_id=}: {denial_reasons=}"
    return PermissionDenied(
        reason=denial_reasons,
        user_message=denial_text,
        log_message=log_text,
    )


@dataclass
class TransferResult:
    shadow_transfer_amount: int
    defending_transfer_amount: int
    transfer_blocked: bool
    transfer_explanation: Optional[str] = None


@sync_to_async
def make_flagged_auth(
    attacking_account: Account,
    defending_account: Account,
    attacker_input: str,
    moderation_metadata: dict,
    llm_output: Optional[str] = None,
) -> None:
    with transaction.atomic():
        if restrictions.check_restriction_enabled(attacking_account, defending_account):
            restriction_id = defending_account.restriction_id
        else:
            restriction_id = None
        auth_attempt = AuthAttempt.objects.create(
            defending_account=defending_account,
            attacking_account=attacking_account,
            was_successful=False,
            attacker_input=attacker_input,
            llm_output=llm_output,
            defense=defending_account.active_defense,
            defense_pre_prompt=defending_account.active_defense.pre_prompt,
            defense_post_prompt=defending_account.active_defense.post_prompt,
            access_code=defending_account.active_defense.access_code,
            llm_choice=defending_account.active_defense.llm_choice,
            attacking_balance=attacking_account.balance,
            defending_balance=defending_account.balance,
            restriction_id=restriction_id,
        )
        Flag.objects.create(
            report_source="automatic",
            authattempt=auth_attempt,
            metadata=moderation_metadata,
        )


def recompute_shadow_balance_inplace(account: Account) -> int:
    # Over the course of SHADOW_BALANCE_INTERVAL, the shadow balance will linearly
    # increase by SHADOW_BALANCE_CAP. (Capped at SHADOW_BALANCE_CAP).
    current_time = timezone.now()
    time_since_last_update = current_time - account.shadow_balance_updated_at
    time_interval_fraction = max(
        0, (time_since_last_update / settings.SHADOW_BALANCE_INTERVAL)
    )
    shadow_balance_increment = time_interval_fraction * settings.SHADOW_BALANCE_CAP
    account.shadow_balance = min(
        account.shadow_balance + shadow_balance_increment,
        settings.SHADOW_BALANCE_CAP,
    )
    account.shadow_balance_updated_at = current_time
    return int(account.shadow_balance)


@dataclass(frozen=True)
class TransferCharges:
    steal_amount: int
    shadow_transfer: int


def get_transfer_charges(
    *, attacking_account: Account, defending_account: Account
) -> TransferCharges:
    """Generates transfer charges.

    Also updates the shadow balance, so this should happen in a transaction
    with the rest of the transfer.
    """
    # Defender gets money from the attacker's shadow balance
    shadow_balance = recompute_shadow_balance_inplace(attacking_account)
    shadow_transfer = int(settings.SHADOW_BALANCE_TRANSFER_FRACTION * shadow_balance)

    # Get the steal amount
    if defending_account.active_defense.post_prompt:
        # Post prompt not empty
        transfer_rate = settings.TRANSFER_CHARGE_FRACTION_WITH_POST_PROMPT
    else:
        # Post prompt empty
        transfer_rate = settings.TRANSFER_CHARGE_FRACTION_WITHOUT_POST_PROMPT
    steal_amount = int(defending_account.balance * transfer_rate)
    return TransferCharges(steal_amount=steal_amount, shadow_transfer=shadow_transfer)


@sync_to_async
def async_make_auth_and_transfer(
    attacking_account: Account,
    defending_account: Account,
    was_successful: bool,
    attacker_input: str,
    llm_output: str,
    moderation_response: ModerationResponse,
    auto_moderated: bool,
) -> TransferResult:
    """Make an AuthAttempt to transfer money between accounts.

    This takes place in a transaction to ensure that it is atomic. It also
    double-checks that the accounts are not already locked, to prevent
    double-attacks.

    If moderation_response says the llm_output is inappropriate,
    then we also create a Flag for the transaction.

    auto_moderated indicates whether all relevant moderation checks succeeded
    (there are actually two—one for the attack text and one for the LLM
    response).

    For an explanation of why we take both moderation_response and
    auto_moderated as parameters, see this GitHub issue:

    https://github.com/qxcv/tensor-trust/pull/171#discussion_r1281128001

    Returns: a TransferResult object containing the amount of money transferred
        (if any), whether the transfer was blocked, and an explanation of what
        happened (to show to the user)."""
    result = TransferResult(
        shadow_transfer_amount=0,
        defending_transfer_amount=0,
        transfer_blocked=False,
        transfer_explanation=None,
    )

    with transaction.atomic():
        # get updated balances & lock states
        attacking_account.refresh_from_db()
        # we need to do more work to refresh the defending account while keeping
        # cooldown info
        defending_account = accounts_with_cooldown_info(attacker=attacking_account).get(
            pk=defending_account.pk
        )

        initial_attacking_balance = attacking_account.balance
        initial_defending_balance = defending_account.balance

        # check whether we're testing our own account
        testing_own_account = (
            attacking_account.account_id == defending_account.account_id
        )

        if restrictions.check_restriction_enabled(attacking_account, defending_account):
            restriction_id = defending_account.restriction_id
        else:
            restriction_id = None

        # Check permissions again, since account balances and lock states may
        # have changed
        match testing_own_account, user_attack_permission(
            attacker=attacking_account, defender=defending_account
        ):
            case _, PermissionDenied(user_message=user_message):
                result.transfer_blocked = True
                result.transfer_explanation = f"Transfer was blocked: {user_message}"
                was_successful = False  # not successful b/c transfer was blocked
            case False, _:
                # We're not testing our own account and permission to attack is
                # still valid, so let's transfer money!
                xfer_charges = get_transfer_charges(
                    attacking_account=attacking_account,
                    defending_account=defending_account,
                )

                if was_successful:
                    if defending_account.active_defense.post_prompt:
                        # Post prompt not empty
                        transfer_explanation_suffix = (
                            f"The defender has a two-part defense, so "
                            f"you steal {int(settings.TRANSFER_CHARGE_FRACTION_WITH_POST_PROMPT * 100)}% of "
                            f"their money!"
                        )
                    else:
                        # Post prompt empty
                        transfer_explanation_suffix = (
                            f"The defender has a one-part defense, so "
                            f"you only steal {int(settings.TRANSFER_CHARGE_FRACTION_WITHOUT_POST_PROMPT * 100)}% of "
                            f"their money!"
                        )
                    result.defending_transfer_amount = xfer_charges.steal_amount

                    # set unlock_at time to the future
                    defending_account.unlock_at = get_unlock_time()

                    # The account was hacked. If it was hacked with a restriction, update the restriction
                    if restriction_id:
                        # This restriction wasn't too hard, so we can increase the tier.
                        higher_tier = restrictions.get_higher_difficulty(
                            defending_account
                        )
                        defending_account.restriction_id = (
                            restrictions.get_random_restriction_id(higher_tier)
                        )
                        defending_account.restriction_id_updated_at = timezone.now()

                    result.transfer_explanation = (
                        f"💰 You stole ${result.defending_transfer_amount}. "
                    )
                    result.transfer_explanation += transfer_explanation_suffix
                else:
                    # On unsuccessful attacks, we don't take money from the defender,
                    # and we give the shadow balance to the defender.
                    result.transfer_explanation = None
                    result.shadow_transfer_amount = xfer_charges.shadow_transfer

                defending_account.balance += result.shadow_transfer_amount
                attacking_account.shadow_balance -= result.shadow_transfer_amount

                defending_account.balance -= result.defending_transfer_amount
                attacking_account.balance += result.defending_transfer_amount

        auth_attempt = AuthAttempt.objects.create(
            defending_account=defending_account,
            attacking_account=attacking_account,
            timestamp=timezone.now(),
            was_successful=was_successful,
            attacker_input=attacker_input,
            llm_output=llm_output,
            attacking_transfer_amount=result.shadow_transfer_amount,
            defending_transfer_amount=result.defending_transfer_amount,
            defense=defending_account.active_defense,
            defense_pre_prompt=defending_account.active_defense.pre_prompt,
            defense_post_prompt=defending_account.active_defense.post_prompt,
            access_code=defending_account.active_defense.access_code,
            llm_choice=defending_account.active_defense.llm_choice,
            auto_moderated=auto_moderated,
            attacking_balance=initial_attacking_balance,
            defending_balance=initial_defending_balance,
            restriction_id=restriction_id,
        )
        match moderation_response:
            case ModerationResponseInappropriate(
                failure_message=_, endpoint_metadata=moderation_metadata
            ):
                Flag.objects.create(
                    report_source="automatic",
                    authattempt=auth_attempt,
                    metadata=moderation_metadata,
                )

        # save all
        auth_attempt.save()
        defending_account.save()
        attacking_account.save()

    log.info(f"Attack generated AuthAttempt: {auth_attempt}")

    return result


async def streaming_llm_output(
    attacking_account: Account,
    attack_prompt: str,
    defense_pre_prompt: str,
    defense_post_prompt: str,
    llm_choice: LLMChoice,
    max_tokens: int,
    enable_gpt_repeat_hack_for_defenses: bool,
) -> AsyncIterable[Tuple[str, bool]]:
    username_hash = hashlib.sha256(
        str(attacking_account.account_id).encode("utf8")
    ).hexdigest()
    generic_err_msg = (
        "⚠️ Could not generate AI response; please refresh and try again. "
        "If the error persists, contact robust-llms@berkeley.edu for help."
    )
    try:
        match rate_limit_result := await rate_limit_check(
            account_id=attacking_account.account_id,
            bucket_configs=settings.API_RATE_LIMIT_BUCKETS,
        ):
            case RateLimitResult(request_allowed=False, wait_time_s=wait_time_s):
                raise RateLimitException(wait_time_s=wait_time_s)
            case RateLimitResult(request_allowed=True, wait_time_s=_):
                pass  # explicit pass
            case _:
                raise RuntimeError(f"Unexpected rate limit result: {rate_limit_result}")
        completion_chunks = await call_llm_api(
            user=f"username-sha256-{username_hash}",
            llm_input=LLMInput(
                pre_prompt=defense_pre_prompt,
                attack_or_access_code=attack_prompt,
                post_prompt=defense_post_prompt,
            ),
            max_tokens=max_tokens,
            llm_choice=llm_choice,
            enable_gpt_repeat_hack_for_defenses=enable_gpt_repeat_hack_for_defenses,
        )
        try:
            while True:
                next_text = await asyncio.wait_for(
                    completion_chunks.__anext__(), timeout=settings.API_TIMEOUT_SECONDS
                )
                yield next_text, True

        except StopAsyncIteration:
            pass
        except asyncio.TimeoutError as ex:
            # We only catch this here so we can log it (to differentiate between failures here
            # and failures in the original chat.completion.create call.)
            log.info(
                f"LLM response timed out during streaming for {attacking_account.account_id=}; {ex}"
            )
            raise ex
        finally:
            await completion_chunks.aclose()

    # FIXME(sam): also need to handle PaLM and Claude errors
    # (probably we should do this at a lower level)
    except (OpenAIError, asyncio.TimeoutError) as ex:
        log.warning(
            "Timeout or OpenAIError while generating LLM response "
            f"{attacking_account.account_id=}: {type(ex)} {ex}"
        )
        # handle the weird error we started getting for some prompts on 2023-11-17
        if (
            isinstance(ex, OpenAIError)
            and ex.user_message == INSCRUTABLE_MYSTERY_OPENAI_ERROR
        ):
            yield (
                "⚠ The AI provider rejected your input prompt. Please try "
                "again with a different prompt.",
                False,
            )
        else:
            yield (generic_err_msg, False)
    except RateLimitException as ex:
        log.info(
            f"User was rate-limited while querying LLM and asked to wait "
            f"{ex.wait_time_s:.4g}s ({attacking_account.account_id=}): {ex}"
        )
        ret_msg = (
            "⚠️ LLM Querying failed due to rate limit; please try again in "
            f"{math.ceil(ex.wait_time_s)} seconds."
        )
        if ex.wait_time_s > 10:
            # add an explanation if we ask them to wait for ages
            ret_msg += (
                " This is to slow down bots. Please contact us on Discord if you "
                "are hitting this rate limit from normal (non-bot) play."
            )
        yield (
            ret_msg,
            False,
        )
    except Exception as ex:
        log.error(
            "Unhandled error while generating LLM response "
            f"{attacking_account.account_id=}: {type(ex)} {ex}"
        )
        yield (generic_err_msg, False)


def validate_attacker_input(
    attacker_input: str,
    attacker_account: Account,
    defender_account: Account,
) -> Optional[str]:
    if (
        attacker_input is not None
        and not get_number_tokens(
            attacker_input, settings.MAX_ATTACK_LENGTH_TOKENS
        ).verification_passed
    ):
        log.info(
            f"User attempted to log in with attack token length over {settings.MAX_ATTACK_LENGTH_TOKENS}"
        )
        return f"Attack failed because the attack string is too long (>{settings.MAX_ATTACK_LENGTH_TOKENS} tokens)"

    if restrictions.check_restriction_enabled(attacker_account, defender_account):
        restriction = restrictions.get_restriction(defender_account)
        restriction_result = restriction.is_valid(attacker_input)
        if not restriction_result.is_valid:
            log.info(
                f"Attack failed because it violates the restriction. {restriction.id}, {attacker_input}, {restriction_result.error_message}"
            )
            return f"Restriction violation: {restriction_result.error_message}"


def validate_sandbox_input(
    *,
    defense_pre_prompt: str,
    attacker_input: str,
    defense_post_prompt: str,
) -> Optional[str]:
    parts = [
        (
            defense_pre_prompt,
            "opening defense",
            settings.MAX_DEFENSE_PRE_PROMPT_LENGTH_TOKENS,
        ),
        (attacker_input, "attack", settings.MAX_ATTACK_LENGTH_TOKENS),
        (
            defense_post_prompt,
            "closing defense",
            settings.MAX_DEFENSE_POST_PROMPT_LENGTH_TOKENS,
        ),
    ]
    for part, part_name, token_limit in parts:
        if not get_number_tokens(part, token_limit).verification_passed:
            log.info(
                f"Sandbox attack failed, {part_name} is too long (>{token_limit} tokens)"
            )
            return f"Attack failed because the {part_name} is too long (>{token_limit} tokens)"


def refresh_restriction_if_old(account: Account) -> None:
    """If it has been at least 24 hours since the restriction was updated,
    then refresh the restriction to a new one."""
    if account.restriction_id_updated_at < timezone.now() - timedelta(hours=24):
        # 24 hrs without the restriction updated means the restriction is too hard.
        # Move down a difficulty level.
        account.restriction_id = restrictions.get_random_restriction_id(
            restrictions.get_lower_difficulty(account)
        )
        account.restriction_id_updated_at = timezone.now()
        account.save()
