from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple, TypeAlias

import pandas as pd
from ttdata.common import ACCESS_GRANTED_REGEX, is_access_granted_medium

# for indexing into authattempts
ATTACKER_ID_COL = "attacking_account_id"
DEFENDER_ID_COL = "defending_account_id"


@dataclass(frozen=True)
class TrajectoryKey:
    attacker_id: int
    defender_id: int


AttackTrajectories: TypeAlias = Dict[TrajectoryKey, List[int]]


def attack_trajectories(attacks: pd.DataFrame) -> AttackTrajectories:
    """Compute 'attack trajectories': sequences of attacks made by one attacker
    on one defender, ordered by time."""
    assert attacks.index.name == "id"
    attack_columns_subset = attacks[[ATTACKER_ID_COL, DEFENDER_ID_COL, "timestamp"]]

    trajectories: Dict[TrajectoryKey, List[Tuple[datetime, int]]] = {}
    # iterate through rows of the dataframe
    for attack_id, row in attack_columns_subset.iterrows():
        attacker_id = int(row[ATTACKER_ID_COL])
        defender_id = int(row[DEFENDER_ID_COL])
        key = TrajectoryKey(attacker_id=attacker_id, defender_id=defender_id)
        if key not in trajectories:
            trajectories[key] = []
        timestamp = row["timestamp"]
        # keeps type checkers happy
        assert isinstance(timestamp, datetime), f"{timestamp=} and {type(timestamp)=}"
        assert isinstance(attack_id, int), f"{attack_id=} and {type(attack_id)=}"
        trajectories[key].append((timestamp, attack_id))

    # sort by time (with ID as secondary sort)
    final_trajectories = {}
    for key in list(trajectories.keys()):
        trajectories[key].sort()
        final_trajectories[key] = [attack_id for _, attack_id in trajectories[key]]
        del trajectories[key]

    return final_trajectories


def add_forward_back_links(
    attacks: pd.DataFrame, attack_trajectories: AttackTrajectories
) -> pd.DataFrame:
    """Compute 'forward links' (ID of next attack by the attacker on the
    defender) and 'back links' (ID of previous attack by the attacker on the
    defender). Returns a dataframe with columns 'traj_next' and 'traj_prev'
    containing the ID of the next and previous trajectory, respectively. The
    frame is indexed by ID, like `attacks`."""
    # make sure that 'id' is the index of attacks
    assert attacks.index.name == "id", f"{attacks.index.name=}"

    # Initialize a DataFrame with 'traj_next' and 'traj_prev' columns (initially missing NaN)
    link_data = pd.DataFrame(
        index=attacks.index, columns=["traj_next", "traj_prev"], dtype=pd.Int64Dtype()
    )

    for traj_attack_ids in attack_trajectories.values():
        for fst, snd in zip(traj_attack_ids, traj_attack_ids[1:]):
            link_data.at[fst, "traj_next"] = snd
            link_data.at[snd, "traj_prev"] = fst

    # you can use attacks.join(add_forward_back_links(attacks,
    # attack_trajectories)) to combine the original frame and the returned frame
    return link_data


def attack_is_access_code(attacks: pd.DataFrame) -> pd.Series:
    """Create a feature column indicating whether the attack was approximately
    equal to the access code."""
    rv = attacks["attacker_input"].str.lower() == attacks["access_code"].str.lower()
    rv.name = "attack_is_access_code"
    return rv


def response_contains_access_code(attacks: pd.DataFrame) -> pd.Series:
    """Create a feature column indicating whether the LLM response contained the
    access code (modulo case)."""

    def _check(row: pd.Series) -> bool:
        llm_output_anycase: Optional[str] = row["llm_output"]
        if llm_output_anycase is None:
            return False  # failure, cannot contain access code
        access_code: str = row["access_code"].lower()
        llm_output = llm_output_anycase.lower()
        return access_code in llm_output

    rv = attacks.apply(_check, axis=1)
    rv.name = "response_contains_access_code"
    return rv


def output_is_access_granted_strong(attacks: pd.DataFrame) -> pd.Series:
    """Creates a new feature indicating whether the access granted regex matches."""

    def says_access_granted(s: Optional[str]) -> bool:
        return s is not None and ACCESS_GRANTED_REGEX.match(s) is not None

    llm_outputs = attacks["llm_output"]
    assert isinstance(llm_outputs, pd.Series)  # to calm down type checker
    rv = llm_outputs.map(says_access_granted)
    # Old name:
    # rv.name = "output_says_access_granted"
    rv.name = "output_is_access_granted_strong"
    return rv


def accounts_used_against(attacks: pd.DataFrame) -> pd.Series:
    """How many accounts has each attack been used against?"""
    used_against = defaultdict(set)
    for _, attack in attacks.iterrows():
        attacker_input = attack["attacker_input"]
        target = attack["defending_account_id"]
        used_against[attacker_input].add(target)

    def _annotate(attacker_input: str) -> int:
        return len(used_against[attacker_input])

    new_series = attacks["attacker_input"].apply(_annotate)
    new_series.name = "accounts_attack_used_against"
    return new_series


def output_is_access_granted_medium(attacks: pd.DataFrame) -> pd.Series:
    """Feature column indicating whether the attack succeeded according to the medium-strength criterion."""

    def _annotate(llm_output: Optional[str]) -> bool:
        if llm_output is None:
            return False
        return is_access_granted_medium(llm_output)

    new_series = attacks["llm_output"].apply(_annotate)
    new_series.name = "output_is_access_granted_medium"
    return new_series


def success_not_access_code(attacks: pd.DataFrame) -> pd.Series:
    """Has the attack even been successfully used against an account in a situation where it was not the access code?"""
    successful_non_access_codes = (
        ~attacks["attack_is_access_code"] & attacks["output_is_access_granted_medium"]
    )
    success_inputs = set(attacks["attacker_input"][successful_non_access_codes])
    rv = attacks["attacker_input"].isin(success_inputs)
    rv.name = "attack_ever_granted_access_without_access_code"
    return rv


def is_first_pass_hijacking_candidate(attacks: pd.DataFrame) -> pd.Series:
    """Is this attacker input a first-pass candidate for hijacking? Needs to satisfy two criteria:
    - Attacker input must have been used successfully against at least one account
      in a situation where it was not the access code.
    - Attacker input must have been used against at least two different accounts in total.
    - There's more than one word in the input"""
    rv = (
        attacks["attack_ever_granted_access_without_access_code"]
        & (attacks["accounts_attack_used_against"] >= 2)
        & (
            (attacks["attacker_input"].str.split().str.len() > 1)
            | (attacks["attacker_input"].str.len() > 30)
        )
    )
    rv.name = "is_first_pass_hijacking_candidate"
    return rv


def prompt_extraction_chain_heuristic(attacks: pd.DataFrame) -> pd.Series:
    """Create an prompt_extraction_chain_heuristic column that is true iff the
    attack is a user's last unsuccessful attack on another player before they
    successfully entered the access code."""
    not_self_attack = attacks[ATTACKER_ID_COL] != attacks[DEFENDER_ID_COL]
    says_access_granted = attacks["output_is_access_granted_strong"]
    attack_is_access_code = attacks["attack_is_access_code"]
    # get IDs for succesful other-attacks where the attacker just entered the
    # access code
    starting_points = attacks[
        not_self_attack & says_access_granted & attack_is_access_code
    ].index
    # we'll backchain using prev_attack_ids
    prev_attack_ids = attacks["traj_prev"]

    def backchain_to_find_prompt_extraction(
        start_attack_id: int,
    ) -> Any:
        """Find the previous unsuccessful attack before start_attack_id."""
        attack_id = prev_attack_ids.loc[start_attack_id]
        while not pd.isnull(attack_id):
            if (
                not says_access_granted.loc[attack_id]
                and not attack_is_access_code.loc[attack_id]
            ):
                return attack_id
            attack_id = prev_attack_ids.loc[attack_id]
        return pd.NA  # we coldn't find an unsuccessful past attack

    # backchain from the starting points
    extraction_attack_ids = (
        starting_points.map(backchain_to_find_prompt_extraction).dropna().unique()
    )

    # create a boolean column
    # frst convert the index to a series so that we can use isin
    assert attacks.index.name == "id"
    id_series = pd.Series(attacks.index, index=attacks.index)
    rv = id_series.isin(extraction_attack_ids)
    rv.name = "prompt_extraction_chain_heuristic"
    assert isinstance(rv, pd.Series)  # to keep type checker happy

    return rv


# Old hijacking heuristic, deprecated 2023-09-24 in favor of a stricter
# heursitic + adversarial filtering
def OLD_may_be_hijacking(attacks: pd.DataFrame) -> pd.Series:
    """Heuristic column indicating that the attack might be a hijack.

    This is much simpler than prompt extraction. We assume an attack is
    hijacking if:

    - It's not a self-attack
    - It's not the access code
    - It's successful

    We don't need to explicitly exclude prompt extractions because the
    prompt_extraction_chain_heuristic column cannot be true unless the attack is
    unsuccessful."""
    not_self_attack = attacks[ATTACKER_ID_COL] != attacks[DEFENDER_ID_COL]
    says_access_granted = attacks["output_is_access_granted_strong"]
    attack_is_access_code = attacks["attack_is_access_code"]
    rv = not_self_attack & ~attack_is_access_code & says_access_granted
    rv.name = "OLD_may_be_hijacking"
    assert isinstance(rv, pd.Series)  # to keep type checker happy
    return rv


def add_extra_feature_columns(attacks: pd.DataFrame) -> pd.DataFrame:
    # add all the extra feature columns defined in this file

    print("Making doubly-linked list of attack 'trajectories'")
    trajectories = attack_trajectories(attacks)
    forward_back_links = add_forward_back_links(attacks, trajectories)
    attacks = attacks.join(forward_back_links)

    print("Adding 'attack_is_access_code' column")
    attacks = attacks.join(attack_is_access_code(attacks))

    print("Adding 'response_contains_access_code' column")
    attacks = attacks.join(response_contains_access_code(attacks))

    print("Adding 'output_says_access_granted' column")
    attacks = attacks.join(output_is_access_granted_strong(attacks))

    print("Adding 'prompt_extraction_chain_heuristic' column")
    attacks = attacks.join(prompt_extraction_chain_heuristic(attacks))

    # Old column, no longer used as of 2023-09-24
    print("Adding 'OLD_may_be_hijacking' column")
    attacks = attacks.join(OLD_may_be_hijacking(attacks))

    print("Adding 'accounts_used_against' column")
    attacks = attacks.join(accounts_used_against(attacks))

    print("Adding 'output_is_access_granted_medium' column")
    attacks = attacks.join(output_is_access_granted_medium(attacks))

    print("Adding 'attack_ever_granted_access_without_access_code' column")
    attacks = attacks.join(success_not_access_code(attacks))

    print("Adding 'is_first_pass_hijacking_candidate' column")
    attacks = attacks.join(is_first_pass_hijacking_candidate(attacks))

    return attacks
