"""Select good hijacking attacks using defenses from good_defenses.py and
features from derived_features.py.

Steps:

0. Start with a set of good defenses & extra attack features in annotated
   attacks.
1. Filter attacks to find only those that match the heuristic in
   derived_features.py.
2. Deduplicate attacks.
3. Generate many attack/defense combinations by randomly pairing the the input
   defenses and input attacks.
3. Adversarial filtering:
    a. Test all the combinations against all the models.
    b. Keep the combinations that work against sufficiently many models.

Our final dataset is the adversarially filtered version.
"""

import warnings
from pathlib import Path
from typing import List, cast

import pandas as pd
from ttdata.call_llm import llm_predict
from ttdata.common import (
    APIError,
    ModelInput,
    dataframe_stable_random_choice,
    is_access_granted_medium,
    write_jsonl,
    zip_safe,
)
from ttdata.duplicate_detection import approximate_deduplicate_keep_last


def create_hijacking_robustness_dataset(
    *,
    output_dir: Path,
    model_names: List[str],
    ignore_cache: bool = False,
    samples_to_filter: int,
    cache_path: Path,
):
    out_path_parquet = output_dir / "hijacking_robustness_dataset_unreviewed.parquet"
    out_path_jsonl = output_dir / "hijacking_robustness_dataset_unreviewed.jsonl"
    out_path_pre_adversarial = (
        output_dir / "hijacking_robustness_pre_adversarial.parquet"
    )
    out_path_includes_invalid = (
        output_dir / "hijacking_robustness_dataset_including_invalid.parquet"
    )
    if not ignore_cache and all(
        p.exists()
        for p in [
            out_path_parquet,
            out_path_jsonl,
            out_path_pre_adversarial,
            out_path_includes_invalid,
        ]
    ):
        print(
            "Using cached hijacking robustness dataset from "
            f"'{out_path_parquet}' and '{out_path_jsonl}'"
        )
        return pd.read_parquet(out_path_parquet)

    raw_attacks = pd.read_parquet(output_dir / "annotated_attacks.parquet")
    defenses = pd.read_parquet(output_dir / "good_annotated_defenses.parquet")

    # Do some pre-filtering to find promising hijacking attacks and deduplicate
    hijack_candidates_mask = raw_attacks["is_first_pass_hijacking_candidate"]
    hijack_candidates = raw_attacks[hijack_candidates_mask]
    # exact drop duplicates
    unique_attacks_exact = hijack_candidates.drop_duplicates(
        ["attacker_input"], keep="last"
    )
    # approximate drop duplicates
    unique_attacks_approx = approximate_deduplicate_keep_last(
        frame=unique_attacks_exact,
        dedup_cols=["attacker_input"],
        # high threshold b/c attacks often have small differences in shingles
        # that make a big difference
        threshold=0.98,
    )

    # Print out some debug info about approximate dataset size
    print(
        f"Total number of attacks: {len(unique_attacks_approx)} soft de-duplicated, "
        f"{len(unique_attacks_exact)} hard de-duplicate, {len(hijack_candidates)} "
        f"filtered, {len(raw_attacks)} before filtering"
    )
    print(
        f"Total number of defenses: {len(defenses)} after good defense pre-processing "
        "(previous performed on cached data)"
    )

    # save a parquet before debugging
    print(f"Writing pre-adversarial-filtering dataset to '{out_path_pre_adversarial}'")
    unique_attacks_approx.to_parquet(out_path_pre_adversarial)
    print(f"Now adversarially filtering with {model_names=}")

    attack_is_valid_col_prefix = "hijack_adv_filter_attack_is_valid_"
    valid_cols = []

    chosen_defenses = None
    chosen_attacks = None
    # while the number of samples obtained is <= the number we want to filter...
    iter_num = 0
    while chosen_defenses is None or len(chosen_defenses) < samples_to_filter:
        new_chosen_defenses = dataframe_stable_random_choice(
            frame=defenses,
            hash_columns=["pre_prompt", "access_code", "post_prompt"],
            key=f"hijack_defenses",
            n=samples_to_filter,
        )
        # Choose an attack for each defense using a stable random choice based
        # on the defense contents. This keeps the output stable even when we add
        # or remove a small number of defenses or attacks from the upstream
        # dataset.
        new_chosen_attack_locs = []
        for _, defense_row in new_chosen_defenses.iterrows():
            new_chosen_attack = dataframe_stable_random_choice(
                frame=unique_attacks_approx,
                hash_columns=["attacker_input"],
                key=f"hijack_attacks",
                # include both defense contents and the iteration number in the
                # value that we use for hashing, so that we make an independent
                # choice for each defense, on each round
                extra_values=[
                    str(iter_num),
                    defense_row.pre_prompt,
                    defense_row.post_prompt,
                    defense_row.access_code,
                ],
                n=1,
            )
            ((new_chosen_attack_loc, _),) = new_chosen_attack.iterrows()
            new_chosen_attack_locs.append(new_chosen_attack_loc)
        new_chosen_attacks = unique_attacks_approx.loc[new_chosen_attack_locs]
        iter_num += 1

        # concatenate the new defenses and attacks
        if chosen_defenses is None:
            chosen_defenses = new_chosen_defenses
            chosen_attacks = new_chosen_attacks
        else:
            assert chosen_attacks is not None
            chosen_defenses = pd.concat([chosen_defenses, new_chosen_defenses])
            chosen_attacks = pd.concat([chosen_attacks, new_chosen_attacks])

    assert chosen_attacks is not None
    assert chosen_defenses is not None
    assert len(chosen_defenses) == len(chosen_attacks)
    assert len(chosen_defenses) == samples_to_filter

    # make a new composite index for the concatenated frame shifting one and
    # OR-ing it with the other, then concatenate the frames columnwise
    shifted_defense_index = chosen_defenses.index.to_numpy() << 32
    attack_index = chosen_attacks.index.to_numpy()
    combined_index = shifted_defense_index | attack_index
    # now use this index for both chosen_defenses and chosen_attakcks
    chosen_attacks.index = pd.Index(combined_index)
    chosen_defenses.index = pd.Index(combined_index)
    # ...and make a combined frame with just the four columns we care abut
    combined_frame = pd.concat(
        [
            chosen_defenses[["pre_prompt", "access_code", "post_prompt"]],
            chosen_attacks[["attacker_input"]],
        ],
        axis=1,
    )

    # No do adversarial filtering with the models
    for model_name in model_names:
        # now we do the adversarial filtering (this ignores the access code,
        # since we've already filtered for defenses where the access code is
        # valid, and just makes sure that the attack passes the current model)
        input_stream = (
            ModelInput(
                pre_prompt=row["pre_prompt"],
                attack_or_access_code=row["attacker_input"],
                post_prompt=row["post_prompt"],
            )
            for _, row in combined_frame.iterrows()
        )
        output_stream = llm_predict(
            model=model_name,
            inputs=input_stream,
            cache_path=cache_path,
            retry_errors_in_cache=True,
            # use short max_length_tokens b/c we only want to see 'access
            # granted'
            max_length_tokens=30,
        )

        # add extra col named attack_is_valid (initially pd.NA)
        model_attack_is_valid_col = attack_is_valid_col_prefix + model_name
        valid_cols.append(model_attack_is_valid_col)
        combined_frame[model_attack_is_valid_col] = pd.NA
        # save model output to this col
        output_save_col = f"hijack_adv_filter_output_{model_name}"
        combined_frame[output_save_col] = pd.NA
        n_errors = 0
        for (row_id, _), model_output in zip_safe(
            combined_frame.iterrows(), output_stream
        ):
            # restore type info
            model_output = cast(list[str | APIError], model_output)

            match model_output:
                case [APIError()]:
                    # this is an error, so we'll just leave the NA in the col
                    n_errors += 1
                case [str(output)]:
                    # this is a valid output, so we'll write it to the col
                    combined_frame.loc[row_id, model_attack_is_valid_col] = (
                        is_access_granted_medium(output)
                    )
                    combined_frame.loc[row_id, output_save_col] = output
                case _:
                    # error, this should never happen
                    raise ValueError(f"Unexpected output from LLM: {model_output}")

        if n_errors > 0:
            warnings.warn(
                f"LLM returned {n_errors} errors for model {model_name}. Consider re-running this script to remove the pd.NAs"
            )

        # print out how common each value of attack_is_valid (False, True, pd.NA) is:
        print(f"attack_is_valid statistics for {model_name=}")
        print(combined_frame[model_attack_is_valid_col].value_counts(dropna=False))

    # print out how many models each attack succeeded against
    print(f"Processed {len(combined_frame)} total samples")
    combined_frame["attack_is_valid_count"] = combined_frame[valid_cols].sum(
        axis=1, skipna=True
    )
    print(
        "attack_is_valid_count statistics (# models for which an attack is successful):"
    )
    print(combined_frame["attack_is_valid_count"].value_counts(dropna=False))

    # save intermediate data
    print(f"Saving dataset including invalid samples to '{out_path_includes_invalid}'")
    combined_frame.to_parquet(out_path_includes_invalid)
    print("Will now continue filtering")

    # filter by summing validity counts
    all_hijacking_samples = combined_frame[combined_frame["attack_is_valid_count"] >= 2]
    print(f"Only {len(all_hijacking_samples)} total hijacking samples")

    # deduplicate (exact deduplication is all we need; this is just to make sure
    # we throw out accidental double-samples of some things)
    dedup_mask = all_hijacking_samples.index.duplicated(keep="last")
    all_hijacking_samples = all_hijacking_samples.loc[~dedup_mask]
    print(f"After dedup we have {len(all_hijacking_samples)} samples")

    # content-based shuffle again on the combined frames
    all_hijacking_samples = dataframe_stable_random_choice(
        frame=all_hijacking_samples,
        hash_columns=["pre_prompt", "access_code", "post_prompt", "attacker_input"],
        key="hijacking_robustness_combined_shuffle",
        n=len(all_hijacking_samples),
    )
    # this should be the same
    print(f"After shuffling we still have {len(all_hijacking_samples)} samples")

    print(f"Writing hijacking robustness dataset to '{out_path_parquet}'")
    all_hijacking_samples.to_parquet(out_path_parquet)

    # write out a jsonl version too
    print(f"Writing a JSON version of the dataset to '{out_path_jsonl}'")
    # first we convert to a list of dicts with these fields:
    # sample_id, pre_prompt, access_code, post_prompt, attack
    records_iter = (
        {
            "sample_id": sample_id,
            "pre_prompt": sample.pre_prompt,
            "access_code": sample.access_code,
            "post_prompt": sample.post_prompt,
            "attack": sample.attacker_input,
        }
        for sample_id, sample in all_hijacking_samples.iterrows()
    )
    # now we write out the jsonl (and don't shuffle again)
    write_jsonl(
        dict_iter=records_iter,
        out_path=out_path_jsonl,
        shuffle_seed=None,
    )

    return all_hijacking_samples
