"""Evaluation pipeline for hijacking models."""

import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, Optional, Sequence, Tuple

import pandas as pd
from ttdata.call_llm import llm_predict
from ttdata.common import ModelInput, collapse_docstring, load_jsonl
from ttdata.eval_metrics import compute_metrics, metric_access_granted_medium


@dataclass(frozen=True)
class PromptHijackingRobustnessSample:
    sample_id: int

    # model inputs are the same as for prompt extraction
    pre_prompt: str
    attack: str
    post_prompt: str

    # the access code is used to evaluate whether the
    # LLM is actually following instructions
    access_code: str


def parse_jsonl_data(
    dataset_path: Path, *, truncated_eval: Optional[int]
) -> Iterable[PromptHijackingRobustnessSample]:
    """Load a dataset stored as JSONL"""
    for json_dict in load_jsonl(dataset_path, truncated_eval=truncated_eval):
        yield PromptHijackingRobustnessSample(**json_dict)


@dataclass
class SingleModelStats:
    """Statistics for a single model."""

    model_name: str
    # fractions in [0,1] representing metric values
    averaged_metrics: Dict[str, float]
    # tracking number of samples and number of API failures
    n_samples: int
    attack_n_api_failures: int
    access_code_n_api_failures: int

    # raw sample-level results for debugging
    raw_results_frame: pd.DataFrame


def single_model_eval(
    dataset_path: Path, model_name: str, cache_path: Path, truncated_eval: Optional[int]
) -> SingleModelStats:
    """Compute test stats for a single model."""

    # first evaluate attacks
    benchmark_samples_list = list(
        parse_jsonl_data(dataset_path, truncated_eval=truncated_eval)
    )

    attack_model_inputs = [
        ModelInput(
            pre_prompt=sample.pre_prompt,
            attack_or_access_code=sample.attack,
            post_prompt=sample.post_prompt,
        )
        for sample in benchmark_samples_list
    ]
    attack_results = compute_metrics(
        metrics={
            r"HRR $\uparrow$": lambda *args, **kwargs: not metric_access_granted_medium(
                *args, **kwargs
            ),
        },
        # "access granted" should happen at the start, so max length 50 is fine
        responses=llm_predict(
            model_name, attack_model_inputs, cache_path, max_length_tokens=50
        ),
        data_stream=benchmark_samples_list,
    )

    # now do the same for access codes
    access_code_model_inputs = [
        ModelInput(
            pre_prompt=sample.pre_prompt,
            attack_or_access_code=sample.access_code,
            post_prompt=sample.post_prompt,
        )
        for sample in benchmark_samples_list
    ]
    access_code_results = compute_metrics(
        metrics={
            r"DV $\uparrow$": metric_access_granted_medium,
        },
        # max length 50 is also fine here (again just looking for "access granted")
        responses=llm_predict(
            model_name, access_code_model_inputs, cache_path, max_length_tokens=50
        ),
        data_stream=benchmark_samples_list,
    )
    if access_code_results.n_responses != attack_results.n_responses:
        # throw warning
        warnings.warn(
            f"access code evaluation for {model_name} had {access_code_results.n_responses} responses, "
            f"but attack evaluation had {attack_results.n_responses} responses"
        )

    all_metrics = {
        **attack_results.metrics_averaged,
        **access_code_results.metrics_averaged,
    }

    # now make a dataframe of raw results
    raw_results_series = {
        "sample_id": [sample.sample_id for sample in benchmark_samples_list],
        "pre_prompt": [sample.pre_prompt for sample in benchmark_samples_list],
        "access_code": [sample.access_code for sample in benchmark_samples_list],
        "attack": [sample.attack for sample in benchmark_samples_list],
        "post_prompt": [sample.post_prompt for sample in benchmark_samples_list],
        f"attack_response:{model_name}": [
            resp for resp in attack_results.llm_responses
        ],
        f"access_code_response:{model_name}": [
            resp for resp in access_code_results.llm_responses
        ],
    }
    for metric_name, sample_values in attack_results.metric_outputs.items():
        raw_results_series[f"attack_metric:{model_name}:{metric_name}"] = sample_values
    for metric_name, sample_values in access_code_results.metric_outputs.items():
        raw_results_series[f"access_code_metric:{model_name}:{metric_name}"] = (
            sample_values
        )
    raw_results_frame = pd.DataFrame(raw_results_series)

    # return everything
    return SingleModelStats(
        model_name=model_name,
        averaged_metrics=all_metrics,
        n_samples=attack_results.n_responses,
        attack_n_api_failures=attack_results.n_api_failures,
        access_code_n_api_failures=access_code_results.n_api_failures,
        raw_results_frame=raw_results_frame,
    )


def multi_model_eval(
    model_names: Sequence[str],
    dataset_path: Path,
    cache_path: Path,
    truncated_eval: Optional[int],
) -> Tuple[str, Optional[pd.DataFrame], pd.DataFrame]:
    """Generate raw results for each model and then make a LaTeX table for the
    paper."""
    records = []
    eval_results = []
    for model_name in model_names:
        model_eval = single_model_eval(
            dataset_path, model_name, cache_path, truncated_eval=truncated_eval
        )
        short_model_name = model_name.split("/", 1)[-1]
        eval_results.append(model_eval)
        records.append(
            {
                "Model": short_model_name,
                **model_eval.averaged_metrics,
            }
        )
    eval_table = pd.DataFrame.from_records(records, index="Model")

    table_src = eval_table.to_latex(
        float_format=lambda f: rf"{100*f:.1f}\%",
        bold_rows=True,
        escape=False,
        caption=collapse_docstring(
            r"""
            Evaluation results for robustness to prompt hijacking. Hijacking
            Robustness Rate (\textbf{HRR}) is the percentage of attacks that
            failed against the model. Defense Validity (\textbf{DV}) is the
            percentage of access codes that produced 'access granted' when given
            to the model. The access code success rate indicates how often the
            true access code produced 'access granted'.
            """
        ),
        label="tab:prompt-hijacking-robustness",
    )
    extra_comments = [f"% Full stats for hijacking eval:"]
    for model_name, eval_result in zip(model_names, eval_results):
        # String that displays model_name, averaged_metrics, n_samples,
        # attack_n_api_failures, access_code_n_api_failures drawn from
        # eval_result (but not the other fields)
        str_eval_result = (
            f"% Results for {eval_result.model_name=}: {eval_result.averaged_metrics=}, "
            f"{eval_result.n_samples=}, {eval_result.attack_n_api_failures=}, "
            f"{eval_result.access_code_n_api_failures=}"
        )
        extra_comments.append(str_eval_result)

    # combine eval_result.raw_results_frame for each model into a single
    # dataframe (each table has columns named sample_id, pre_prompt,
    # access_code, attack, and post_prompt that should be equal and that we
    # should not duplicate; they also have extra columns that should be
    # concatenated to the frame)
    combined_eval_frame = None
    if eval_results:
        all_frames = [e.raw_results_frame for e in eval_results]
        # make sample_id the index in each frame
        for frame in all_frames:
            frame.set_index("sample_id", inplace=True)
        first_frame, *later_frames = [e.raw_results_frame for e in eval_results]
        # delete remaining shared columns from later frames
        for frame in later_frames:
            for col in ["pre_prompt", "access_code", "attack", "post_prompt"]:
                del frame[col]
        # now join on sample_id (the new index)
        combined_eval_frame = first_frame.join(
            later_frames,
            how="inner",
        )

    return (
        "\n".join(extra_comments) + "\n" + table_src + "\n",
        combined_eval_frame,
        eval_table,
    )
