import argparse
import copy
import json
import os
import re
from typing import List

import random
from tqdm import tqdm
import torch

from lib.head_mask_inference import load_model  # noqa: E402


def _load_json(path: str):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _load_text(path: str) -> str:
    with open(path, "r", encoding="utf-8") as f:
        return f.read()


def _build_messages(system_prompt: str, question: str) -> List[dict]:
    return [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": question},
    ]


def _prepare_batch(system_prompt: str, items: List[dict], tokenizer):
    prompts = []
    for d_item in items:
        question = d_item.get("instruction", "")
        messages = _build_messages(system_prompt, question)
        prompts.append(tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True))
    return tokenizer(prompts, return_tensors="pt", padding=True, truncation=True)


def _generate_batch(model, tok, encoded, max_new_tokens):
    if encoded["input_ids"].numel() == 0:
        return []
    encoded = encoded.to(model.device)
    out = model.generate(
        **encoded,
        max_new_tokens=max_new_tokens,
        do_sample=False,
        eos_token_id=tok.eos_token_id,
        pad_token_id=tok.pad_token_id,
    )
    outputs = []
    padding_side = getattr(tok, "padding_side", "right")
    for i, row in enumerate(out):
        if padding_side == "left":
            prompt_len = encoded["input_ids"].shape[1]
        else:
            prompt_len = int(encoded["attention_mask"][i].sum().item())
        gen_ids = row.tolist()
        outputs.append(tok.decode(gen_ids[prompt_len:], skip_special_tokens=True))
    return outputs


def _normalize(text: str) -> str:
    text = text.lower()
    text = re.sub(r"[^a-z0-9\s]+", "", text)
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def _check_ans(model_answer, ans):
    modelans = _normalize(model_answer)
    if "||" in ans:
        for a in ans.split("||"):
            an = _normalize(a)
            if an in modelans:
                return True
        return False
    if "&&" in ans:
        for a in ans.split("&&"):
            an = _normalize(a)
            if an not in modelans:
                return False
        return True
    ans = _normalize(ans)
    return ans in modelans


@torch.inference_mode()
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--victim_model_path", required=True)
    parser.add_argument("--data_path", required=True)
    parser.add_argument("--data_path2", default=None)
    parser.add_argument("--victim_system_path", required=True)
    parser.add_argument("--output_path", required=True)
    parser.add_argument("--output_path2", default=None)
    parser.add_argument("--batch_size", type=int, default=6)
    parser.add_argument("--data_size", type=int, default=-1)
    parser.add_argument("--max_new_tokens", type=int, default=256)
    args = parser.parse_args()

    random.seed(42)
    torch.manual_seed(42)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(42)

    data = _load_json(args.data_path)
    data2 = _load_json(args.data_path2) if args.data_path2 else None
    system_prompt = _load_text(args.victim_system_path)

    model, tok = load_model(args.victim_model_path)
    if getattr(tok, "padding_side", None) != "left":
        tok.padding_side = "left"

    if data2 is not None and len(data) != len(data2):
        raise ValueError("Paired datasets must have the same length.")

    if args.data_size > 0:
        data = data[:args.data_size]
        if data2 is not None:
            data2 = data2[:args.data_size]

    kept_items = []
    kept_items2 = [] if data2 is not None else None
    removed_count = 0
    batch_items = []
    batch_items2 = []

    for idx, item in enumerate(tqdm(data)):
        batch_items.append(copy.deepcopy(item))
        if data2 is not None:
            item2 = copy.deepcopy(data2[idx])
            if item.get("instruction") != item2.get("instruction") or item.get("output") != item2.get("output"):
                raise ValueError(f"Paired item mismatch at index {idx}.")
            batch_items2.append(item2)
        if len(batch_items) < args.batch_size:
            continue

        encoded = _prepare_batch(system_prompt, batch_items, tok)
        outputs = _generate_batch(model, tok, encoded, args.max_new_tokens)
        for b_idx, d_item in enumerate(batch_items):
            response = outputs[b_idx]
            if _check_ans(response, d_item.get("output", "")):
                removed_count += 1
            else:
                kept_items.append(d_item)
                if kept_items2 is not None:
                    kept_items2.append(batch_items2[b_idx])
        batch_items = []
        batch_items2 = []

    if batch_items:
        encoded = _prepare_batch(system_prompt, batch_items, tok)
        outputs = _generate_batch(model, tok, encoded, args.max_new_tokens)
        for b_idx, d_item in enumerate(batch_items):
            response = outputs[b_idx]
            if _check_ans(response, d_item.get("output", "")):
                removed_count += 1
            else:
                kept_items.append(d_item)
                if kept_items2 is not None:
                    kept_items2.append(batch_items2[b_idx])

    output_dir = os.path.dirname(args.output_path)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)
    with open(args.output_path, "w", encoding="utf-8") as f:
        json.dump(kept_items, f, indent=2, ensure_ascii=False)
    if kept_items2 is not None:
        if not args.output_path2:
            raise ValueError("output_path2 is required when data_path2 is provided.")
        output_dir2 = os.path.dirname(args.output_path2)
        if output_dir2:
            os.makedirs(output_dir2, exist_ok=True)
        with open(args.output_path2, "w", encoding="utf-8") as f:
            json.dump(kept_items2, f, indent=2, ensure_ascii=False)

    print(f"Saved filtered dataset to: {args.output_path}")
    if args.output_path2:
        print(f"Saved filtered dataset to: {args.output_path2}")
    print(f"Total: {len(data)} | Removed: {removed_count} | Kept: {len(kept_items)}")


if __name__ == "__main__":
    main()
