"""Code for making a dump of the raw dataset, minus naughty stuff that we don't
want in there."""

import datetime
import hashlib
import subprocess
import unicodedata
from pathlib import Path
from typing import Dict, Optional, Tuple

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


def NFD(text):
    return unicodedata.normalize("NFD", text)


def canonical_caseless(text):
    return NFD(NFD(text).casefold())


def canonical_equal(text1, text2):
    return canonical_caseless(text1) == canonical_caseless(text2)


def canonical_contains(text, substring):
    return canonical_caseless(substring) in canonical_caseless(text)


# example of usage
if False:
    assert "å".casefold() != "å".casefold()
    assert canonical_caseless("å") == canonical_caseless("å")


# These all violate ToS in some way (words are de-anonymizing or appear in a
# profane context). Word list was created for the September 10th dump and will
# have to be updated for future dumps.
bad_words = [
    "nigger",
    "alex turner",
    "sam toyer",
    "berkeley",
    "Stuart Russell",
    "Ethan Mendes",
    "BAIR",
    "Justin Svegliato",
    "Luke Bailey",
    "Olivia Watkins",
    "SCUM Manifesto",
    "hitler",
    "anal ",
    "pieter",
    "transgender",
    "pornographic",
    "porn",
    "porno",
    "lezbiyen",
    "horni",
    "penis",
    "blowjob",
    "Fuck bitches",
    "tits",
    "breasts",
    "dick",
    "mom is a hoe",
    "kotyKDisthebest",
    "neega",
]


def is_naughty_string(s: str) -> bool:
    return any(canonical_contains(text=s, substring=bad_word) for bad_word in bad_words)


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


def impl_make_raw_dump(output_dir: Path, blake_key: str) -> None:
    attacks_parquet_path = output_dir / "annotated_attacks.parquet"
    defenses_parquet_path = output_dir / "gameui_defense.parquet"

    attacks_frame = pd.read_parquet(attacks_parquet_path)
    defenses_frame = pd.read_parquet(defenses_parquet_path)

    # Step 1: anonymizing attacker/defender IDs. We make a list of all the
    # attacker/defender IDs, then HMAC them with a secret key. Place in a dict
    # and resolve any collisions.

    # get unique account IDs
    all_acct_ids = pd.concat(
        [
            attacks_frame["attacking_account_id"],
            attacks_frame["defending_account_id"],
            defenses_frame["creator_account_id"],
        ]
    ).unique()
    # assign an anonymous ID to each user by running blake
    anonymous_ids = {}
    anonymous_id_values = set()
    for acct_id in all_acct_ids:
        current_hash = None
        ctr = 0
        while current_hash is None or current_hash in anonymous_id_values:
            # run blake2b but only take 53 bits (i.e. convert to unsigned int
            # and take % 2^53)
            target_bits = 53
            hashed = hashlib.blake2b(
                f"{ctr}-{acct_id}".encode("utf-8"),
                key=blake_key.encode("utf-8"),
                digest_size=8,
            ).digest()
            current_hash = int.from_bytes(hashed, byteorder="big") & (
                (1 << target_bits) - 1
            )
            ctr += 1
        if ctr > 1:
            print(f"Collision resolved for {acct_id} after {ctr} tries")
        assert current_hash is not None
        anonymous_ids[acct_id] = current_hash

    # Next step: mapping (defender_id, pre_prompt, access_code, post_prompt) to
    # defense_ids and timestamps for the defenses.  There may be duplicates, if
    # someone saved the same defense multiple times, which is why we save the
    # dupes in a dict.
    defense_ids: Dict[Tuple[int, str, str, str], Dict[int, datetime.datetime]] = {}
    for defense_id, defense_row in defenses_frame.iterrows():
        key = (
            defense_row["creator_account_id"],
            defense_row["pre_prompt"],
            defense_row["access_code"],
            defense_row["post_prompt"],
        )
        if key not in defense_ids:
            defense_ids[key] = {}
        assert isinstance(defense_id, int), f"{defense_id=}, {type(defense_id)=}"
        defense_ids[key][defense_id] = defense_row["timestamp"]

    def find_defense_for_attack(attack_row: pd.Series) -> Optional[int]:
        """Find the right defense ID for the attack, if possible.

        First we look for exact matches on defender ID, pre-prompt, access code,
        and post-prompt.
        If we get more than one, we try to find the last defense created before
        the attack.
        If they all have a timestamp _after_ the attack, we give up."""
        key = (
            attack_row["defending_account_id"],
            attack_row["defense_pre_prompt"],
            attack_row["access_code"],
            attack_row["defense_post_prompt"],
        )
        if key not in defense_ids:
            return None
        defense_id_to_timestamp = defense_ids[key]
        defense_ids_sorted = sorted(defense_id_to_timestamp.items(), key=lambda x: x[1])
        defense_ids_sorted.reverse()
        for defense_id, timestamp in defense_ids_sorted:
            if timestamp < attack_row["timestamp"]:
                return defense_id
        return None

    n_sanitized = 0
    n_total = 0

    def sanitize_string(s: Optional[str]) -> Optional[str]:
        nonlocal n_sanitized, n_total
        n_total += 1
        if s is None:
            return None
        if is_naughty_string(s):
            n_sanitized += 1
            return None
        return s

    attack_dict_stream = (
        {
            "attack_id": attack_id,
            "attacker_id_anonymized": anonymous_ids[attack_row["attacking_account_id"]],
            "defender_id_anonymized": anonymous_ids[attack_row["defending_account_id"]],
            "defense_id": find_defense_for_attack(attack_row),
            "attacker_balance_before": attack_row["attacking_balance"],
            "defender_balance_before": attack_row["defending_balance"],
            # Naming for the transfer amounts is very weird:
            # 1. 'attacking_transfer_amount' is the amount the _defender_ wins
            #    (higher = more money for defender)
            # 2. 'defending_transfer_amount' is the amount the _attacker_ wins
            #    (higher = more money for attacker)
            # The names below are hopefully easier to understand.
            "attacker_balance_gain": attack_row["defending_transfer_amount"],
            "defender_balance_gain": attack_row["attacking_transfer_amount"],
            "opening_defense": sanitize_string(attack_row["defense_pre_prompt"]),
            "attacker_input": sanitize_string(attack_row["attacker_input"]),
            "closing_defense": sanitize_string(attack_row["defense_post_prompt"]),
            "access_code": sanitize_string(attack_row["access_code"]),
            "llm_choice": attack_row["llm_choice"],
            "llm_output": sanitize_string(attack_row["llm_output"]),
            "output_is_access_granted": is_access_granted_strong(
                attack_row["llm_output"]
            ),
            "is_self_attack": attack_row["attacking_account_id"]
            == attack_row["defending_account_id"],
            "timestamp": attack_row["timestamp"].isoformat(),
        }
        for attack_id, attack_row in attacks_frame.iterrows()
    )

    defense_dict_stream = (
        {
            "defense_id": defense_id,
            "defender_id_anonymized": anonymous_ids[defense_row["creator_account_id"]],
            "opening_defense": sanitize_string(defense_row["pre_prompt"]),
            "closing_defense": sanitize_string(defense_row["post_prompt"]),
            "access_code": sanitize_string(defense_row["access_code"]),
            "llm_choice": defense_row["llm_choice"],
            "llm_output": sanitize_string(defense_row["access_code_output"]),
            "output_is_access_granted": is_access_granted_strong(
                defense_row["access_code_output"]
            ),
            "timestamp": defense_row["timestamp"].isoformat(),
        }
        for defense_id, defense_row in defenses_frame.iterrows()
    )

    # now write them both out as jsonlines
    output_dir.mkdir(parents=True, exist_ok=True)
    attacks_output_path = output_dir / "raw_dump_attacks.jsonl"
    defenses_output_path = output_dir / "raw_dump_defenses.jsonl"
    print(f"Writing attacks to '{attacks_output_path}' (may be slow)")
    write_jsonl(attack_dict_stream, attacks_output_path, shuffle_seed=None)
    print(f"Writing defenses to '{defenses_output_path}' (a bit less slow)")
    write_jsonl(defense_dict_stream, defenses_output_path, shuffle_seed=None)

    print(
        f"Sanitized {n_sanitized} out of {n_total} strings ({n_sanitized / n_total:.2%})"
    )

    print("Writing pbzip2-compressed copies (may be slow if you don't have many cores)")
    subprocess.run(["pbzip2", "-v", "-k", "-f", attacks_output_path], check=True)
    subprocess.run(["pbzip2", "-v", "-k", "-f", defenses_output_path], check=True)
    print("Compressed copies written. We're done!")
