import argparse
from dataclasses import asdict
from datetime import datetime
import gc
import json
from pathlib import Path
import pickle

import numpy as np
import torch
from transformers import set_seed

import softopt

from api import EvalInstance
from judge import Judge, JudgeConfig
from softopt import SoftOptConfig
from utils import (
    generate, 
    handle_non_serializable, 
    load_model_and_tokenizer,
)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-m", "--model_name_or_path", type=str, default=None)
    parser.add_argument("-j", "--judge_name_or_path", type=str, default="cais/HarmBench-Llama-2-13b-cls")
    parser.add_argument("--benchmark", type=str, default=None)
    parser.add_argument("--start", type=int, default=0)
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output_dir", type=str, default=None)
    parser.add_argument("--save_norms", action="store_true")
    # RepE arguments
    parser.add_argument("--use_repe", action="store_true")
    parser.add_argument("--layer_id", type=list, default=list(range(-11, -21, -1)))
    parser.add_argument("--repe_coeff", type=float, default=1.0)
    # SoftOpt (Input Embedding Attack) arguments
    parser.add_argument("--num_steps", type=int, default=500)
    parser.add_argument("--optim_str_init", type=str, default="x x x x x x x x x x x x x x x x x x x x")
    parser.add_argument("--run_softopt", action="store_true")
    parser.add_argument("--lr", type=float, default=0.0001)
    parser.add_argument("--early_stop_loss", type=float, default=0.05)
    parser.add_argument("--log_soft_embeds", action="store_true")
    parser.add_argument("--log_losses", action="store_true")
    # Generation arguments
    parser.add_argument("--temperature", type=float, default=None)
    parser.add_argument("--do_sample", type=bool, default=False)
    parser.add_argument("--top_p", type=float, default=1.0)
    parser.add_argument("--max_new_tokens", type=int, default=512)
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--prefill", type=bool, default=False)
    parser.add_argument("--use_template", type=bool, default=True)
    # Judge arguments
    parser.add_argument("--judge_dtype", type=str, default="bfloat16")
    parser.add_argument("--judge_tensor_parallel_size", type=int, default=1)
    parser.add_argument("--judge_temperature", type=int, default=0)
    parser.add_argument("--judge_max_new_tokens", type=int, default=1)

    args = parser.parse_args()
    return args

def evaluate(
    model_name_or_path,
    benchmark, 
    start,
    limit,
    seed,
    repe_config,
    softopt_config,
    gen_kwargs,
    judge_config
):
    if seed is not None:
        set_seed(seed)
        torch.use_deterministic_algorithms(True, warn_only=True)

    model, tokenizer = load_model_and_tokenizer(model_name_or_path)
    if repe_config is not None:
        init_rep_control(model, tokenizer, layer_id=repe_config["layer_id"], repe_coeff=repe_config["repe_coeff"])
    
    with open(benchmark) as f:
        benchmark = json.load(f)
        instances = []
        for d in benchmark["data"]:
            instance = EvalInstance(**d)
            if instance.context is not None:
                instance.messages[-1]["content"] = instance.context + "\n\n---\n\n" + instance.messages[-1]["content"]
            instances.append(instance)

        instances = instances[start:]
        if limit is not None:
            instances = instances[:limit]
    
    # Generate completions from the model
    if softopt_config is None:
        generate(model, tokenizer, instances, gen_kwargs)
    else: # softopt requires a special generation
        gen_kwargs.pop("batch_size"); gen_kwargs.pop("compute_norms"); gen_kwargs.pop("prefill"); gen_kwargs.pop("use_template")
        log_soft_embeds = softopt_config.pop("log_soft_embeds")
        log_losses = softopt_config.pop("log_losses")
        softopt_config = SoftOptConfig(**softopt_config)

        assert instances[0].default_target is not None, "SoftOpt requires target strings"
        for instance in instances:
            target = instance.default_target + ":\n\n"
            softopt_result = softopt.run(model, tokenizer, instance.behavior, target, softopt_config)
            input_embeds = softopt_result.input_embeds.to("cuda")
            with torch.no_grad():
                output = model.generate(
                    inputs_embeds=input_embeds, 
                    pad_token_id=tokenizer.pad_token_id,
                    return_dict_in_generate=True,
                    **gen_kwargs,
                )
                sequences = output.sequences.detach().cpu()
                generation = tokenizer.decode(sequences[0], skip_special_tokens=True)
                
                del output
                gc.collect()
                torch.cuda.empty_cache()

            instance.generation = generation
            if log_soft_embeds:
                instance.input_embeds = input_embeds.tolist()
            if log_losses:
                instance.losses = softopt_result.losses

    print(f"Sample input: {instances[0].messages}")
    print(f"Sample generation: {instances[0].generation}")

    del model
    gc.collect()
    torch.cuda.empty_cache()

    # Run evaluation using LLM judge
    judge_config = JudgeConfig(**judge_config)
    judge = Judge(judge_config)
    score = judge.evaluate(instances)

    return score, instances

def main():
    args = parse_args()
    args_dict = vars(args)
    for k, v in args_dict.items():
        print(f"{k}: {v}")

    if args.use_repe:
        repe_config = {
            "layer_id": args.layer_id,
            "repe_coeff": args.repe_coeff,
        }
    else:
        repe_config = None

    if args.run_softopt:
        softopt_config = {
            "num_steps": args.num_steps,
            "optim_str_init": args.optim_str_init,
            "lr": args.lr,
            "early_stop_loss": args.early_stop_loss,
            "log_soft_embeds": args.log_soft_embeds,
            "log_losses": args.log_losses,
            "seed": args.seed,
        }
    else:
        softopt_config = None
    
    gen_kwargs = {
        "temperature": args.temperature,
        "do_sample": args.do_sample,
        "top_p": args.top_p,
        "max_new_tokens": args.max_new_tokens,
        "batch_size": args.batch_size,
        "compute_norms": args.save_norms,
        "prefill": args.prefill,
        "use_template": args.use_template,
    }

    judge_config = {
        "model_name_or_path": args.judge_name_or_path,
        "dtype": args.judge_dtype,
        "tensor_parallel_size": args.judge_tensor_parallel_size,
        "temperature": args.judge_temperature,
        "max_new_tokens": args.judge_max_new_tokens,
        "seed": args.seed,
    }

    score, instances = evaluate(
        args.model_name_or_path,
        args.benchmark,
        args.start,
        args.limit,
        args.seed,
        repe_config,
        softopt_config,
        gen_kwargs,
        judge_config
    )

    print(f"Evaluation score: {score}")

    if args.output_dir is not None:
        output_dir = Path(args.output_dir)
        if not output_dir.is_dir():
            output_dir.mkdir(parents=True, exist_ok=True)
        
        samples = []
        activation_norms = []
        generated_tokens = []
        for instance in instances:
            sample = asdict(instance)
            activation_norms.append(instance.activation_norms)
            generated_tokens.append(instance.tokens)
            samples.append({k:v for k, v in sample.items() if k not in ("activation_norms, tokens")})
        
        result = {"args": args_dict, "time": datetime.now(), "score": score, "samples": samples}
        dumped = json.dumps(result, indent=2, default=handle_non_serializable, ensure_ascii=False)
        with open(output_dir.joinpath("log.json"), "w") as f:
            f.write(dumped)

        if args.save_norms:
            with open(output_dir.joinpath("activation_norms.pkl"), "wb") as f:
                pickle.dump(activation_norms, f)

            with open(output_dir.joinpath("tokens.pkl"), "wb") as f:
                pickle.dump(generated_tokens, f)
        
if __name__ == "__main__":
    main()