#!/usr/bin/env python3
"""Generate IDIRAK evaluation outputs for native-speaker review."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import sacrebleu
import torch
from datasets import load_dataset
from peft import AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-id", default="Qwen/Qwen3-4B")
    parser.add_argument("--adapter-id")
    parser.add_argument(
        "--dataset-id", default="UyghurAI/idirak-uyghur-instructions"
    )
    parser.add_argument("--split", default="test")
    parser.add_argument("--output-file", default="evaluations/predictions.jsonl")
    parser.add_argument("--max-examples", type=int, default=0)
    parser.add_argument("--max-new-tokens", type=int, default=384)
    parser.add_argument("--no-quantization", action="store_true")
    return parser.parse_args()


def generation_text(tokenizer, messages) -> str:
    return tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )


def main() -> None:
    args = parse_args()
    if not torch.cuda.is_available():
        raise SystemExit("Evaluation requires a CUDA GPU for the configured models")

    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    quantization_config = None
    if not args.no_quantization:
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=dtype,
        )

    load_id = args.adapter_id or args.model_id
    tokenizer = AutoTokenizer.from_pretrained(load_id, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model_class = AutoPeftModelForCausalLM if args.adapter_id else AutoModelForCausalLM
    model = model_class.from_pretrained(
        load_id,
        device_map="auto",
        torch_dtype=dtype,
        quantization_config=quantization_config,
    )
    model.eval()

    dataset = load_dataset(args.dataset_id, split=args.split)
    if args.max_examples > 0:
        dataset = dataset.select(range(min(args.max_examples, len(dataset))))

    output_path = Path(args.output_file).expanduser().resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    translation_predictions = []
    translation_references = []

    with output_path.open("w", encoding="utf-8") as output:
        for row in dataset:
            messages = row["messages"]
            if not messages or messages[-1].get("role") != "assistant":
                raise ValueError(f"{row.get('id', 'unknown')} has no final assistant answer")

            prompt_messages = messages[:-1]
            reference = messages[-1]["content"]
            prompt = generation_text(tokenizer, prompt_messages)
            encoded = tokenizer(
                prompt,
                return_tensors="pt",
                truncation=True,
                max_length=4096,
            ).to(model.device)
            with torch.inference_mode():
                generated = model.generate(
                    **encoded,
                    max_new_tokens=args.max_new_tokens,
                    do_sample=False,
                    repetition_penalty=1.05,
                    pad_token_id=tokenizer.eos_token_id,
                )
            answer_tokens = generated[0][encoded.input_ids.shape[1] :]
            prediction = tokenizer.decode(
                answer_tokens, skip_special_tokens=True
            ).strip()

            record = {
                "id": row.get("id"),
                "category": row.get("category"),
                "prompt_messages": prompt_messages,
                "reference": reference,
                "prediction": prediction,
                "review": {
                    "uyghur_fluency_1_to_5": None,
                    "correctness_1_to_5": None,
                    "instruction_following_1_to_5": None,
                    "safety_1_to_5": None,
                    "overall_1_to_5": None,
                    "notes": "",
                },
            }
            output.write(json.dumps(record, ensure_ascii=False) + "\n")

            if row.get("category") == "translation":
                translation_predictions.append(prediction)
                translation_references.append(reference)

    print(f"Wrote {len(dataset)} predictions to {output_path}")
    if translation_predictions:
        score = sacrebleu.corpus_chrf(
            translation_predictions, [translation_references]
        ).score
        print(f"Translation chrF: {score:.2f}")
    print("Complete the review fields with native Uyghur speakers before release.")


if __name__ == "__main__":
    main()