#!/usr/bin/env python3
"""Inference script for Faruq IndoLaw Alignment Models."""

import argparse
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer


def main() -> None:
    parser = argparse.ArgumentParser(description="Run inference with Faruq IndoLaw models from Hugging Face.")
    parser.add_argument("--repo-id", default="ruwwww/faruq-alignment-models", help="Hugging Face repo ID")
    parser.add_argument("--subfolder", default="models/sft-v2-5epoch", help="Model subfolder within repository")
    parser.add_argument("--base-model", default="Qwen/Qwen2.5-1.5B", help="Base model ID")
    parser.add_argument("--max-new-tokens", type=int, default=512, help="Maximum generated tokens")
    args = parser.parse_args()

    print(f"Loading tokenizer & adapter from {args.repo_id} [{args.subfolder}]...")
    tokenizer = AutoTokenizer.from_pretrained(args.repo_id, subfolder=args.subfolder)
    base_model = AutoModelForCausalLM.from_pretrained(
        args.base_model,
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else None,
    )
    model = PeftModel.from_pretrained(base_model, args.repo_id, subfolder=args.subfolder)
    model.eval()

    messages = [
        {
            "role": "system",
            "content": (
                "Anda adalah asisten hukum Indonesia. Gunakan hanya dakwaan dan fakta persidangan yang diberikan. "
                "Susun pertimbangan hukum yang ringkas dan amar putusan secara setia pada sumber."
            ),
        },
        {
            "role": "user",
            "content": (
                "[SURAT DAKWAAN]\n"
                "Bahwa terdakwa didakwa melanggar Pasal 362 KUHP tentang Pencurian karena telah mengambil satu unit sepeda motor milik korban tanpa izin.\n\n"
                "[FAKTA PERSIDANGAN & ALAT BUKTI]\n"
                "Saksi menerangkan melihat terdakwa mendorong sepeda motor korban di malam hari. "
                "Barang bukti berupa 1 unit sepeda motor dan kunci kontak diamankan. Terdakwa mengakui perbuatannya di depan persidangan."
            ),
        },
    ]

    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    print("\nGenerating court ratio decidendi and dictum...\n")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=args.max_new_tokens,
            do_sample=False,
            repetition_penalty=1.1,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
        )

    completion = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    print("=== MODEL OUTPUT ===")
    print(completion)


if __name__ == "__main__":
    main()