"""FoodGuard binary classification using the checkpoint's original policy."""
import argparse
import json
import re
from pathlib import Path

import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM, AutoTokenizer


class FoodGuard:
    def __init__(self, model, device='auto', revision=None):
        local = Path(model).is_dir()
        common = {'local_files_only': local, 'trust_remote_code': False}
        if revision is not None:
            common['revision'] = revision
        policy_file = Path(model) / 'classifier_policy.txt' if local else Path(
            hf_hub_download(model, 'classifier_policy.txt', revision=revision))
        self.policy = policy_file.read_text(encoding='utf-8').rstrip('\n')
        self.tokenizer = AutoTokenizer.from_pretrained(model, **common)
        use_gpu = device == 'cuda' or (device == 'auto' and torch.cuda.is_available())
        if device == 'cuda' and not torch.cuda.is_available():
            raise RuntimeError('CUDA requested but unavailable')
        self.device = 'cuda:0' if use_gpu else 'cpu'
        dtype = (torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16) if use_gpu else torch.float32
        self.model = AutoModelForCausalLM.from_pretrained(
            model, dtype=dtype, device_map={'': self.device},
            attn_implementation='sdpa', low_cpu_mem_usage=True, **common,
        ).eval()

    def classify(self, text):
        prompt = self.tokenizer.apply_chat_template([
            {'role': 'system', 'content': self.policy},
            {'role': 'user', 'content': text},
        ], tokenize=False, add_generation_prompt=True)
        inputs = self.tokenizer(prompt, add_special_tokens=False, return_tensors='pt')
        length = inputs.input_ids.shape[1]
        if length > 5120:
            raise ValueError(f'Input has {length} tokens; maximum is 5120. Input was not truncated.')
        inputs = inputs.to(self.device)
        with torch.inference_mode():
            output = self.model.generate(
                **inputs, do_sample=False, max_new_tokens=16,
                pad_token_id=self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id,
            )
        answer = self.tokenizer.decode(output[0, length:], skip_special_tokens=True).strip()
        match = re.fullmatch(r'Safety Assessment:\s*(safe|unsafe)', answer, re.I)
        return {'predicted_label': match.group(1).lower() if match else None,
                'final_answer': answer, 'input_tokens': length}


def main():
    parser = argparse.ArgumentParser(description='Classify food-related input with FoodGuard')
    parser.add_argument('--model', required=True, help='Local model directory or Hugging Face model ID')
    parser.add_argument('--revision', help='Optional Hugging Face commit or tag')
    parser.add_argument('--device', choices=['auto', 'cuda', 'cpu'], default='auto')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--text')
    group.add_argument('--input-file', type=Path, help='JSONL containing an input field')
    parser.add_argument('--output-file', type=Path)
    args = parser.parse_args()
    if args.input_file and not args.output_file:
        parser.error('--input-file requires --output-file')
    if args.input_file and args.input_file.resolve() == args.output_file.resolve():
        parser.error('Input and output files must differ')
    guard = FoodGuard(args.model, args.device, args.revision)
    if args.text is not None:
        print(json.dumps(guard.classify(args.text), ensure_ascii=False))
        return
    args.output_file.parent.mkdir(parents=True, exist_ok=True)
    with args.input_file.open(encoding='utf-8') as source, args.output_file.open('w', encoding='utf-8') as target:
        for line_number, line in enumerate(source, 1):
            if not line.strip():
                continue
            row = json.loads(line)
            if not isinstance(row.get('input'), str):
                raise ValueError(f'Line {line_number}: input must be a string')
            record = {'line_number': line_number, **guard.classify(row['input'])}
            for field in ('group_id', 'template_id', 'label'):
                if field in row:
                    record[field] = row[field]
            target.write(json.dumps(record, ensure_ascii=False) + '\n')
            target.flush()


if __name__ == '__main__':
    main()