#!/usr/bin/env python3
import datetime
import random
import socket
import string
import subprocess
import sys
import warnings
from contextlib import contextmanager
from pathlib import Path
from typing import Generator, List, Optional, TypedDict

import click
import matplotlib.pyplot as plt
import pandas as pd
import ttdata.eval_extraction as extraction
import ttdata.eval_hijacking as hijacking
import ttdata.eval_prompt_extraction_detection as prompt_extraction_detection
from adjustText import adjust_text
from ttdata import sample_annotator
from ttdata.call_llm import MODELS_DICT
from ttdata.common import LLMPromptType, load_jsonl, write_jsonl
from ttdata.create_extraction_robustness_dataset import (
    create_extraction_robustness_dataset,
)
from ttdata.create_hijacking_robustness_dataset import (
    create_hijacking_robustness_dataset,
)
from ttdata.derived_features import add_extra_feature_columns
from ttdata.good_defenses import annotate_defenses_with_models, choose_good_defenses
from ttdata.make_raw_dump import impl_make_raw_dump


def random_string(length: int = 10) -> str:
    """Generate a random string of the given length."""
    letters = string.ascii_letters + string.digits
    return "".join(random.choice(letters) for i in range(length))


def get_free_port() -> int:
    """Get a free port from the OS and return it."""
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind(("", 0))
    s.listen(1)
    port = s.getsockname()[1]
    s.close()
    return port


@contextmanager
def temporary_postgres_container(sql_dump_path: Path) -> Generator[str, None, None]:
    """Context manager for a temporary Postgres container."""
    postgres_port = get_free_port()
    postgres_password = random_string(20)
    postgres_container_name = f"temp_postgres_container_{random_string(5)}"
    POSTGRES_DB = "mydatabase"

    print(
        f"Spinning up Postgres container (port={postgres_port}, name={postgres_container_name})"
    )
    subprocess.run(
        [
            "docker",
            "run",
            "--name",
            postgres_container_name,
            "--rm",
            "-e",
            f"POSTGRES_PASSWORD={postgres_password}",
            "-e",
            f"POSTGRES_DB={POSTGRES_DB}",
            "-p",
            f"{postgres_port}:5432",
            # this is a 'temporary anonymous volume" (gets cleaned up on stop)
            "-v",
            "/var/lib/postgresql/data",
            "-d",
            "postgres",
        ],
        check=True,
    )

    try:
        # loop wait on pg_isready with docker exec
        max_wait_time_s = 30
        print(f"Waiting for Postgres to spin up (max {max_wait_time_s}s)")

        subprocess.run(
            [
                "docker",
                "exec",
                "-i",
                postgres_container_name,
                "timeout",
                f"{max_wait_time_s}s",
                "bash",
                "-c",
                f"until pg_isready; do sleep 5; done",
            ],
            check=True,
        )

        # this pulls the dump into the container and runs it, minus the GRANTs
        # (which reference IAM users that don't exist outside Cloud Run)
        safe_sql_dump_path = str(sql_dump_path)
        safe_command = (
            f"< {safe_sql_dump_path} grep -Ev '^GRANT ' "
            f"| docker exec -i {postgres_container_name} "
            f"psql -U postgres -d {POSTGRES_DB}"
        )
        subprocess.run(safe_command, shell=True, check=True)

        # let other code do something with the DB
        yield f"postgresql+psycopg://postgres:{postgres_password}@localhost:{postgres_port}/{POSTGRES_DB}"
    finally:
        print("Deleting Postgres container")
        # docker stop should kill the thing and clean up the volume because of
        # the --rm/-v combo passed to docker run
        subprocess.run(["docker", "stop", postgres_container_name], check=True)
        print("Postgres container deleted")


def process_postgres_data(
    pg_db_uri: str, output_dir: Path, cutoff_date: Optional[datetime.datetime]
) -> None:
    """Process data in Postgres and output to Parquet files."""
    output_dir.mkdir(parents=True, exist_ok=True)
    # read to SQL and write to Parquet
    for table_name in ["gameui_authattempt", "gameui_defense", "gameui_flag"]:
        df = pd.read_sql_query(f"SELECT * FROM {table_name}", pg_db_uri, index_col="id")

        # fancy filtering
        if cutoff_date is not None:
            # convert cutoff date to 11:59PM AOE timestamp, add AOE timezone
            cutoff_date_aoe = cutoff_date.replace(
                hour=23,
                minute=59,
                second=59,
                microsecond=999999,
                tzinfo=datetime.timezone(datetime.timedelta(hours=-12)),
            )
            include_mask = df["timestamp"] <= cutoff_date_aoe
            print(
                f"Filtering {table_name} to include only data up to and including {cutoff_date}"
            )
            # print num discarded and % discarded
            print(
                f"This discards {len(df) - include_mask.sum():,} rows of {len(df):,} ({(1 - include_mask.mean()) * 100:.2f}%)"
            )
            df = df[include_mask]

        df.to_parquet(output_dir / f"{table_name}.parquet")


def do_initial_parquet_load(
    sql_dump_path: Path,
    output_dir: Path,
    ignore_cache: bool,
    cutoff_date: Optional[datetime.datetime],
) -> None:
    output_files = [
        output_dir / f"{table_name}.parquet"
        for table_name in ["gameui_authattempt", "gameui_defense", "gameui_flag"]
    ]
    if all(output_file.exists() for output_file in output_files) and not ignore_cache:
        # cached
        print("Skipping initial load from SQL, will use cached Parquet files")
        return
    if cutoff_date is not None:
        print(f"Will only process data up to and including {cutoff_date}")
    with temporary_postgres_container(sql_dump_path) as pg_db_uri:
        process_postgres_data(pg_db_uri, output_dir, cutoff_date=cutoff_date)


def add_attack_features(output_dir: Path, ignore_cache: bool) -> None:
    # write extra feature columns to attacks
    attacks_path_parquet = output_dir / "annotated_attacks.parquet"
    if attacks_path_parquet.exists() and not ignore_cache:
        # cached
        print("Skipping extra attack features, will use cached Parquet files")
        return

    # read from Parquet and write to Parquet
    attacks = pd.read_parquet(output_dir / "gameui_authattempt.parquet")
    attacks = add_extra_feature_columns(attacks)
    print(f"Writing annotated attacks to '{attacks_path_parquet}'")
    attacks.to_parquet(attacks_path_parquet)


def save_plot(df: pd.DataFrame, output_dir: Path, plot_type="extraction") -> None:
    if plot_type == "extraction":
        y_title = "Extraction Robustness Rate"
        save_name = "extraction_robustness.pdf"
        robustness_col_name = r"ERR $\uparrow$"
    elif plot_type == "hijacking":
        y_title = "Hijacking Robustness Rate"
        save_name = "hijacking_robustness.pdf"
        robustness_col_name = r"HRR $\uparrow$"
    else:
        raise ValueError(f"Unknown plot type {plot_type}")

    model_names = df.index.tolist()
    defense_validity = df[r"DV $\uparrow$"].tolist()
    defense_robustness_rate = df[robustness_col_name].tolist()

    # Turn values into percentages
    defense_validity = [x * 100 for x in defense_validity]
    defense_robustness_rate = [x * 100 for x in defense_robustness_rate]

    # Define model families and their corresponding colors
    family_colors = {
        "gpt": "red",
        "claude": "green",
        "bison": "blue",
        "llama": "purple",
    }

    # Define reference models and their markers
    reference_models = ["gpt-3.5-turbo-0613", "claude-instant-1.2", "chat-bison@001"]
    markers = [
        "o" if model_name not in reference_models else "x" for model_name in model_names
    ]

    # Plotting
    plt.clf()
    plt.figure(figsize=(8, 6))
    plotted_families = []

    for model_name, d_validity, robustness, mark in zip(
        model_names, defense_validity, defense_robustness_rate, markers
    ):
        family = next(f for f in family_colors.keys() if f in model_name.lower())
        label = family if family not in plotted_families else None
        if label:
            plotted_families.append(family)
        plt.scatter(
            d_validity,
            robustness,
            color=family_colors[family],
            marker=mark,
            s=150,
        )

    plt.xlabel("Defense Validity", fontsize=16)
    plt.ylabel(y_title, fontsize=16)
    plt.xlim(0, 100)
    plt.ylim(0, 100)
    plt.grid(True)

    # Family legend with only circles
    handles = [
        plt.Line2D(
            [0],
            [0],
            marker="o",
            color="w",
            label=family,
            markersize=12,
            markerfacecolor=family_colors[family],
        )
        for family in plotted_families
    ]
    leg1 = plt.legend(
        handles=handles,
        title="Model Families",
        title_fontsize=16,
        fontsize=14,
        loc="upper left",
    )

    # Add family legend back to the plot (to allow adding the second legend)
    plt.gca().add_artist(leg1)

    # Legend
    markers_legend = [
        plt.Line2D(
            [0],
            [0],
            marker="o",
            color="w",
            label="Held-Out Model",
            markersize=12,
            markerfacecolor="black",
            linestyle="None",
        ),
        plt.Line2D(
            [0],
            [0],
            marker="x",
            color="black",
            label="Reference Model",
            markersize=12,
            linestyle="None",
        ),
    ]
    leg2 = plt.legend(
        handles=markers_legend,
        title="Model Type",
        title_fontsize=16,
        fontsize=14,
        loc="lower left",
    )
    ax = plt.gca()
    legend_bbox = leg2.get_window_extent().transformed(ax.transData.inverted())
    legend_bbox1 = leg1.get_window_extent().transformed(ax.transData.inverted())

    # Annotations on points
    texts = []
    for i, txt in enumerate(model_names):
        texts.append(
            plt.annotate(
                txt,
                # (min(99.9, defense_validity[i]), defense_robustness_rate[i] - 5),
                (min(99.9, defense_validity[i]), defense_robustness_rate[i]),
                fontsize=14,
                alpha=0.7,
                va="center",
                ha="center",
            )
        )
    adjust_text(
        texts,
        objects=[leg1, leg2],
        time_lim=5.0,
        expand=(1.1, 1.2),
        explode_radius=10.0,
        force_static=(0.2, 0.4),
        avoid_self=False,
        force_text=(0.05, 0.1),
        # objects=[legend_bbox, legend_bbox1],
        # force_pull=(0.01, 0.01),
        # force_text=(0.05, 0.1),
        # force_explode=(0.05, 0.1),
        # add_objects=[legend_bbox, legend_bbox1],
        # lim=50_000,  # default lim=500
        # expand_text=(1.2, 1.5),
        # expand_points=(1.2, 1.5),
        # force_text=(0.1, 0.25),  # default force_text=(0.1, 0.25)
        # force_points=(0.5, 0.5),  # default force_points=(0.2, 0.5)
        # force_objects=(0.1, 0.5),  # default force_objects=(0.1, 0.25),
        # precision=0.001,  # default precision=0.01 (max error 1% of axis range)
        # arrowprops=dict(arrowstyle="-", color='k', lw=0.5),
    )

    plt.tight_layout()
    plt.savefig(output_dir / save_name)
    print(f"Saved plot to {output_dir / save_name}")


def run_hijacking_eval(
    output_dir: Path,
    cache_path: Path,
    truncated_eval: Optional[int],
    model_names: list[str],
) -> None:
    print(f"Running hijacking eval on {len(model_names)} models")
    as_latex, debug_frame, eval_df = hijacking.multi_model_eval(
        model_names,
        output_dir / "hijacking_robustness_dataset.jsonl",
        cache_path,
        truncated_eval=truncated_eval,
    )

    if debug_frame is not None:
        debug_out_path = output_dir / "hijacking_eval_debug.parquet"
        print(f"Writing debug frame for eval to {debug_out_path}")
        debug_frame.to_parquet(debug_out_path)
    else:
        print("No debug frame generated, results are probably empty")

    out_path = output_dir / "hijacking_eval.tex"
    print(f"Writing LaTeX hijacking eval metrics to '{out_path}'")
    with out_path.open("w") as f:
        f.write(as_latex)
    save_plot(eval_df, output_dir, plot_type="hijacking")


def run_extraction_eval(
    output_dir: Path,
    cache_path: Path,
    truncated_eval: Optional[int],
    model_names: list[str],
) -> None:
    print(f"Running extraction eval on {len(model_names)} models")
    as_latex, debug_frame, eval_df = extraction.multi_model_eval(
        model_names,
        output_dir / "extraction_robustness_dataset.jsonl",
        cache_path,
        truncated_eval=truncated_eval,
    )

    if debug_frame is not None:
        debug_out_path = output_dir / "extraction_eval_debug.parquet"
        print(f"Writing debug frame for eval to {debug_out_path}")
        debug_frame.to_parquet(debug_out_path)
    else:
        print("No debug frame generated, results are probably empty")

    out_path = output_dir / "extraction_eval.tex"
    print(f"Writing extraction eval metrics to '{out_path}'")
    with out_path.open("w") as f:
        f.write(as_latex)
    save_plot(eval_df, output_dir, plot_type="extraction")


def run_prompt_extraction_detection_eval(
    *,
    output_dir: Path,
    cache_path: Path,
    truncated_eval: Optional[int],
    model_names: list[str],
    prompt_type: LLMPromptType = LLMPromptType.DETECT_ZERO_SHOT,
) -> None:
    label = "few-shot" if prompt_type == LLMPromptType.DETECT_FEW_SHOT else "zero-shot"

    print(
        f"Running {label} prompt extraction detection eval on {len(model_names)} models"
    )
    as_latex = prompt_extraction_detection.make_eval_table(
        model_names,
        output_dir / "prompt_extraction_detection.jsonl",
        cache_path,
        truncated_eval=truncated_eval,
        prompt_type=prompt_type,
    )

    out_path = output_dir / f"prompt_extraction_detection_eval_{label}.tex"
    print(f"Writing {label} prompt extraction detection eval metrics to '{out_path}'")
    with out_path.open("w") as f:
        f.write(as_latex)


def get_prompt_extraction_detection_dataset(output_dir: Path) -> None:
    print(f"Generating prompt extraction detection dataset")
    attacks = pd.read_parquet(output_dir / "annotated_attacks.parquet")

    # get positives by deduplicating anything marked
    # prompt_extraction_chain_heuristic
    dedup_cols = ["defense_pre_prompt", "attacker_input", "defense_post_prompt"]
    positives = attacks[attacks["prompt_extraction_chain_heuristic"]].drop_duplicates(
        subset=dedup_cols
    )

    print(f"Loaded {len(positives)} positive samples")

    # pick positives from successful extractions
    easy_positives = positives[positives["response_contains_access_code"]]
    hard_positives = positives[~positives["response_contains_access_code"]]

    print(f"Found {len(easy_positives)} easy positives")
    print(f"Found {len(hard_positives)} hard positives")

    # filter outputs that are longer than 20 characters
    # seems to be a good heuristic for getting more meaningful outputs
    hard_positives = hard_positives[
        hard_positives["llm_output"].str.len() > len("Access denied")
    ]

    print(
        f"Found {len(hard_positives)} hard positives after filtering > {len('Access denied')} characters"
    )

    # calculate total balanced lengths for each set:
    # 1. out of positives, should have as many hard as possible, with 30% easy
    # 2. out of negatives, should have 50% hard, 50% easy
    # 2. should have 50% positives, 50% negatives
    # (assumes hard positives are the smallest set)
    total_hard_positives = len(hard_positives)
    total_easy_positives = int((len(hard_positives) * 0.3) / 0.7)
    total_positives = total_hard_positives + total_easy_positives
    total_half_negatives = int(
        (total_positives) / 2
    )  # as we have as many negatives as positives

    def generate_hard_negatives(n_hard_negatives: int = 100) -> pd.DataFrame:
        len_hard_positives = len(hard_positives) - 1
        hard_negatives = pd.DataFrame(columns=["access_code", "llm_output"])

        # generate hard negatives from hard positives by randomizing
        # access codes and outputs
        for _ in range(n_hard_negatives):
            x = random.randint(0, len_hard_positives)
            should_reroll = True
            access_code = hard_positives.iloc[x]["access_code"]

            y = None

            # randomize x, y; if access codes are too similar, pick another y
            while should_reroll:
                y = random.randint(0, len_hard_positives)
                access_code_y = hard_positives.iloc[y]["access_code"]
                should_reroll = (
                    access_code.lower() in access_code_y.lower()
                    or access_code_y.lower() in access_code.lower()
                )

            assert y is not None

            # use final y to get llm_output
            llm_output = hard_positives.iloc[y]["llm_output"]

            # add to hard_negatives
            hard_negatives = pd.concat(
                [
                    hard_negatives,
                    pd.DataFrame(
                        [{"access_code": access_code, "llm_output": llm_output}]
                    ),
                ],
                ignore_index=True,
            )

        return hard_negatives

    # pick negatives
    easy_negatives = attacks[~attacks["may_be_prompt_extraction"]]
    hard_negatives = generate_hard_negatives(total_half_negatives)

    # randomize easy positives & negatives and take the first N
    easy_positives = easy_positives.sample(n=total_easy_positives)
    easy_negatives = easy_negatives.sample(n=total_half_negatives)
    print(
        f"Found {len(hard_negatives)} hard negatives and {len(easy_negatives)} easy negatives"
    )

    # assert our numbers are correct
    assert len(hard_negatives) == total_half_negatives
    assert len(easy_negatives) == total_half_negatives
    assert len(easy_positives) == total_easy_positives
    assert len(hard_positives) == total_hard_positives

    # concat and label positives & negatives
    positives = pd.concat([easy_positives, hard_positives])
    positives["is_prompt_extraction"] = True
    negatives = pd.concat([easy_negatives, hard_negatives])
    negatives["is_prompt_extraction"] = False

    # save positives to parquet
    positives_path = output_dir / "prompt_extraction_detection_positives.parquet"
    print(f"Saving prompt extraction detection positives to {positives_path}")
    positives.to_parquet(positives_path)

    # combine everything and shuffle
    dataset = pd.concat([positives, negatives]).sample(frac=1)
    print(f"Gathered {len(dataset)} total samples")

    # prepare training and validation sets
    cutoff_point = int(len(dataset) * 0.9)
    training_set = dataset[:cutoff_point]
    validation_set = dataset[cutoff_point:]

    # save results to Parquet
    training_path = output_dir / "prompt_extraction_detection_training.parquet"
    validation_path = output_dir / "prompt_extraction_detection_validation.parquet"

    print(f"Saving detection training data to {training_path}")
    training_set.to_parquet(training_path)

    print(f"Saving detection validation data to {validation_path}")
    validation_set.to_parquet(validation_path)

    # save results to JSONL
    extractions_path_jsonl = output_dir / "prompt_extraction_detection.jsonl"
    write_jsonl(
        (
            {
                "sample_id": sample_id,
                "access_code": row["access_code"],
                "llm_output": row["llm_output"],
                "is_prompt_extraction": row["is_prompt_extraction"],
            }
            for sample_id, row in dataset.iterrows()
        ),
        extractions_path_jsonl,
        shuffle_seed=18981,
    )


class PipelineContext(TypedDict):
    """Shared options for the pipeline."""

    ignore_cache: bool
    cache_path: Path
    output_dir: Path


def _set_eval_unlimited(ctx, param, value):
    """Click callback that sets n_eval_samples to use unlimited samples"""
    if value:
        # if 'value' is given then the flag was passed, otherwise it was not
        ctx.params["n_eval_samples"] = None


def _set_eval_limited(ctx, param, value):
    """Click callback that sets n_eval_samples to a limited value."""
    if value is not None:
        # if 'value' is not None then the user actually passed in a numeric
        # limit
        ctx.params["n_eval_samples"] = value
    elif "n_eval_samples" not in ctx.params:
        # otherwise we use the default, unless _set_eval_unlimited has already
        # set it to None
        ctx.params["n_eval_samples"] = 2


@click.group()
@click.option(
    "--ignore-cache",
    is_flag=True,
    help="Ignore the cache and re-run the pipeline from scratch.",
)
@click.option(
    "--cache-path",
    type=click.Path(
        path_type=Path,
        file_okay=True,
        dir_okay=False,
    ),
    default="/home/olivia/lm_cache/lm_cache.db",
    help="Path to the cache database",
)
@click.option(
    "--output-dir",
    type=click.Path(exists=False, path_type=Path, file_okay=False, dir_okay=True),
    help="Directory to write output files.",
    default="./outdata",
    required=True,
)
@click.pass_context
def cli(
    ctx: click.Context,
    ignore_cache: bool,
    cache_path: Path,
    output_dir: Path,
) -> None:
    ctx.ensure_object(dict)
    ctx.obj = PipelineContext(
        ignore_cache=ignore_cache,
        cache_path=cache_path,
        output_dir=output_dir,
    )


@cli.command(
    help="Do initial load from .sql dump to .parquet files",
)
@click.option(
    "--sql-dump-path",
    default="/nas/ucb/sam/lang/tt-private/ro-replica-2024-03-15-prod-db.sql",
    help="path to the raw SQL dump file (defaults to a path on the NAS)",
    type=click.Path(exists=True, path_type=Path, file_okay=True, dir_okay=False),
)
@click.option(
    "--cutoff-date",
    help="only process files up to and including this date (AOE)",
    type=click.DateTime(formats=["%Y-%m-%d"]),
    default=None,
)
@click.pass_context
def initial_load(
    ctx: click.Context, sql_dump_path: Path, cutoff_date: Optional[datetime.datetime]
) -> None:
    pctx: PipelineContext = ctx.obj
    do_initial_parquet_load(
        sql_dump_path,
        pctx["output_dir"],
        ignore_cache=pctx["ignore_cache"],
        cutoff_date=cutoff_date,
    )


@cli.command(
    help="Add extra columns to loaded 'raw' attacks to be used by downstream operations"
)
@click.pass_context
def postprocess_raw_attacks(ctx: click.Context) -> None:
    pctx: PipelineContext = ctx.obj
    add_attack_features(pctx["output_dir"], ignore_cache=pctx["ignore_cache"])


@cli.command(
    help="Run defenses and access codes through a set of models so that we can "
    "later determine which ones give valid output"
)
@click.option(
    "-m",
    "--model",
    "models",
    type=click.Choice(list(MODELS_DICT.keys())),
    multiple=True,
    default=[
        "claude-instant-1.2",
        "gpt-3.5-turbo-0613",
        "chat-bison@001",
    ],
)
@click.pass_context
def check_defense_transfer(ctx: click.Context, models: List[str]) -> None:
    pctx: PipelineContext = ctx.obj
    annotate_defenses_with_models(
        output_dir=pctx["output_dir"],
        model_names=models,
        truncated_eval=None,
        cache_path=pctx["cache_path"],
    )
    print("Done annotating defenses with models, will skip rest of script")


@cli.command(
    help="Use annotated model outputs to choose a set of 'good' defenses for "
    "use in our benchmarks. Also removes defense that fall to trivial attack."
)
@click.pass_context
def generate_good_defenses(ctx: click.Context) -> None:
    pctx: PipelineContext = ctx.obj
    choose_good_defenses(
        output_dir=pctx["output_dir"],
        ignore_cache=pctx["ignore_cache"],
        cache_path=pctx["cache_path"],
    )


@cli.command(
    help="Create a hijacking robustness dataset by identifying potential hijack "
    "attacks and then adversarially filtering to find some that work on at least "
    "one of the given models"
)
@click.option(
    "-m",
    "--model",
    "models",
    type=click.Choice(list(MODELS_DICT.keys())),
    multiple=True,
    default=[
        "claude-instant-1.2",
        "gpt-3.5-turbo-0613",
        "chat-bison@001",
    ],
)
@click.option(
    "--samples-to-filter",
    type=int,
    default=50,
    help="Number of samples to filter for the adversarially filtered "
    "hijacking dataset (this is the number of inputs, not outputs)",
    show_default=True,
)
@click.pass_context
def create_hijacking_dataset(
    ctx: click.Context, models: List[str], samples_to_filter: int
) -> None:
    if samples_to_filter < 5_000:
        warnings.warn(
            f"Testing only {samples_to_filter:,} pairs for the adversarially "
            "filtered hijacking dataset, not the full set of 5,000. Use "
            "--samples-per-model=5000 to use the full set."
        )

    pctx: PipelineContext = ctx.obj
    create_hijacking_robustness_dataset(
        output_dir=pctx["output_dir"],
        model_names=models,
        ignore_cache=pctx["ignore_cache"],
        samples_to_filter=samples_to_filter,
        cache_path=pctx["cache_path"],
    )


@cli.command(
    help="Create an extraction robustness dataset by identifying potential "
    "extraction attacks then adversarially filtering to find combinations "
    "of attack and defense that leak the true access code verbatim when run "
    "through at least one of the reference models"
)
@click.option(
    "-m",
    "--model",
    "models",
    type=click.Choice(list(MODELS_DICT.keys())),
    multiple=True,
    default=[
        "claude-instant-1.2",
        "gpt-3.5-turbo-0613",
        "chat-bison@001",
    ],
)
@click.option(
    "--samples-to-filter",
    type=int,
    default=50,
    help="Number of samples to filter for the adversarially filtered dataset (this is "
    "number of samples we start with; adversarial filtering reduces the number)",
    show_default=True,
)
@click.pass_context
def create_extraction_dataset(
    ctx: click.Context, models: List[str], samples_to_filter: int
) -> None:
    if samples_to_filter < 5_000:
        warnings.warn(
            f"Selecting only {samples_to_filter} to filter for the adversarially "
            "filtered extraction dataset, not the full set of 5,000. Use "
            "--samples-per-model=5000 to use the full set."
        )

    pctx: PipelineContext = ctx.obj
    create_extraction_robustness_dataset(
        output_dir=pctx["output_dir"],
        model_names=models,
        ignore_cache=pctx["ignore_cache"],
        samples_to_filter=samples_to_filter,
        cache_path=pctx["cache_path"],
    )


@cli.command(
    help="Create the prompt extraction detection dataset",
)
@click.pass_context
def create_extraction_detection_dataset(ctx: click.Context) -> None:
    pctx: PipelineContext = ctx.obj
    get_prompt_extraction_detection_dataset(pctx["output_dir"])


@cli.command(
    help="Manually verify samples from either the hijacking or extraction robustness "
    "datasets. Uses a tkinter GUI, so can't run on a server."
)
@click.argument(
    "dataset", type=click.Choice(["hijacking-robustness", "extraction-robustness"])
)
@click.option(
    "--force-show-gui",
    is_flag=True,
    help="Force showing the GUI even if all samples are annotated",
)
@click.pass_context
def verify_samples(
    ctx: click.Context,
    dataset: str,
    force_show_gui: bool,
) -> None:
    pctx: PipelineContext = ctx.obj

    out_dir = pctx["output_dir"]
    match dataset:
        case "hijacking-robustness":
            input_path = out_dir / "hijacking_robustness_dataset_unreviewed.jsonl"
            output_path = out_dir / "hijacking_robustness_dataset.jsonl"
            anno_cache_path = out_dir / "hijacking_robustness_dataset_anno_cache.txt"
        case "extraction-robustness":
            input_path = out_dir / "extraction_robustness_dataset_unreviewed.jsonl"
            output_path = out_dir / "extraction_robustness_dataset.jsonl"
            anno_cache_path = out_dir / "extraction_robustness_dataset_anno_cache.txt"
        case _:
            raise ValueError(f"Unknown dataset {dataset}")

    if output_path.exists() and not pctx["ignore_cache"]:
        pass

    if not input_path.exists():
        raise ValueError(
            f"Input path {input_path} does not exist; probably need to create the "
            "dataset first"
        )

    # first
    def load_and_check_annos():
        anno_results_stringy = sample_annotator.load_results_into_dict(anno_cache_path)
        is_reviewed_by_sample_id = {
            int(sample_id_str): anno_result == "valid"
            for sample_id_str, anno_result in anno_results_stringy.items()
        }
        samples = list(load_jsonl(input_path, truncated_eval=None))
        unannotated_samples = {
            s["sample_id"] for s in samples
        } - is_reviewed_by_sample_id.keys()
        good_samples = [
            s for s in samples if is_reviewed_by_sample_id[s["sample_id"]] is True
        ]
        return unannotated_samples, good_samples

    unannotated_samples, good_samples = load_and_check_annos()
    # show GUI if there are unannotated samples
    if force_show_gui or unannotated_samples:
        print(
            f"Found {len(unannotated_samples)} unannotated samples, showing GUI to annotate the rest"
        )
        reviewer_gui = sample_annotator.SampleReviewer(
            samples_filename=input_path,
            results_filename=anno_cache_path,
        )
        reviewer_gui.mainloop()
        unannotated_samples, good_samples = load_and_check_annos()
    else:
        print(
            "All samples are annotated already, skipping GUI (use --force-show-gui "
            "to make it show up)"
        )

    # don't save if anything was unannotated
    if unannotated_samples:
        print(
            f"WARNING: {len(unannotated_samples)} samples were not annotated, even "
            "after showing GUI. Will not write output file."
        )
        sys.exit(1)

    # save the user-verified samples
    print(f"Writing {len(good_samples)} manually reviewed samples to '{output_path}'")
    write_jsonl(good_samples, output_path, shuffle_seed=None)
    print(f"Done, ready to run eval for {dataset.replace('-', ' ')}")


@cli.command(
    help="Run evals on our extraction and hijacking robustness benchmarks and "
    "produce LaTeX tables. Defaults to running with just one model and a "
    "handful of samples."
)
@click.option(
    "-m",
    "--model",
    "models",
    type=click.Choice(list(MODELS_DICT.keys())),
    multiple=True,
    default=["gpt-3.5-turbo-0613"],
)
@click.option(
    "--truncated-eval",
    type=int,
    help="Truncate evaluation to the first <this many> samples (default 2)",
    callback=_set_eval_limited,
    expose_value=False,
    is_eager=False,
)
@click.option(
    "--no-truncated-eval",
    is_flag=True,
    callback=_set_eval_unlimited,
    expose_value=False,
    is_eager=True,
    help="Eval on full dataset (all available samples)",
)
@click.pass_context
def run_robustness_evals(
    ctx: click.Context, models: List[str], n_eval_samples: Optional[int]
) -> None:
    if n_eval_samples is None:
        print("Running eval on full dataset (all available samples)")
    else:
        print(f"Running truncated LLM eval on first {n_eval_samples} samples")

    pctx: PipelineContext = ctx.obj

    run_hijacking_eval(
        output_dir=pctx["output_dir"],
        cache_path=pctx["cache_path"],
        truncated_eval=n_eval_samples,
        model_names=models,
    )
    run_extraction_eval(
        output_dir=pctx["output_dir"],
        cache_path=pctx["cache_path"],
        truncated_eval=n_eval_samples,
        model_names=models,
    )


@cli.command(
    help="Run extraction detection evals and produce LaTeX tables. "
    "Defaults to running with just one model and a handful of samples."
)
@click.option(
    "-m",
    "--model",
    "models",
    type=click.Choice(list(MODELS_DICT.keys())),
    multiple=True,
    default=["gpt-3.5-turbo-0613"],
)
@click.option(
    "--truncated-eval",
    type=int,
    help="Truncate evaluation to the first <this many> samples (default 2)",
    callback=_set_eval_limited,
    expose_value=False,
    is_eager=False,
)
@click.option(
    "--no-truncated-eval",
    is_flag=True,
    callback=_set_eval_unlimited,
    expose_value=False,
    is_eager=True,
    help="Eval on full dataset (all available samples)",
)
@click.pass_context
def run_detection_evals(
    ctx: click.Context, models: List[str], n_eval_samples: Optional[int]
) -> None:
    if n_eval_samples is None:
        print("Running eval on full dataset (all available samples)")
    else:
        print(f"Running truncated LLM eval on first {n_eval_samples} samples")

    pctx: PipelineContext = ctx.obj

    run_prompt_extraction_detection_eval(
        output_dir=pctx["output_dir"],
        cache_path=pctx["cache_path"],
        truncated_eval=n_eval_samples,
        model_names=models,
        prompt_type=LLMPromptType.DETECT_FEW_SHOT,
    )
    run_prompt_extraction_detection_eval(
        output_dir=pctx["output_dir"],
        cache_path=pctx["cache_path"],
        truncated_eval=n_eval_samples,
        model_names=models,
        prompt_type=LLMPromptType.DETECT_ZERO_SHOT,
    )


@cli.command(
    help="Write raw samples after naughty word removal and anonymization of users. "
    "Uses blake_key to anonymize users (imperfect but makes direct lookup much harder)."
)
@click.argument("blake_key")
@click.pass_context
def make_raw_dump(ctx: click.Context, blake_key: str) -> None:
    pctx: PipelineContext = ctx.obj
    impl_make_raw_dump(pctx["output_dir"], blake_key=blake_key)


def _launch():
    # Sane Click launch script that doesn't catch KeyboardError/EOFError.
    # Copied from here:
    # https://github.com/qxcv/mtil/blob/master/mtil/utils/misc.py#L29
    try:
        with cli.make_context(sys.argv[0], sys.argv[1:]) as ctx:
            cli.invoke(ctx)
    except click.ClickException as e:
        e.show()
        sys.exit(e.exit_code)
    except click.exceptions.Exit as e:
        if e.exit_code == 0:
            sys.exit(e.exit_code)
        raise


if __name__ == "__main__":
    _launch()
