"""Generate the measured selective-restriction adapter's structured response."""

import json
from pathlib import Path
import torch
from huggingface_hub import hf_hub_download
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "itsalissonsilva/Restriction-Learning-Experiment"
BASE = "google/gemma-3-1b-it"
BASE_REV = "dcc83ea841ab6100d6b47a070329e1ba4cf78752"


class RestrictionModel:
    def __init__(self, revision=None, device="cuda"):
        report = json.loads(
            Path(hf_hub_download(REPO, "restriction_results.json", revision=revision)).read_text()
        )
        self.policy, self.format = report["policy"], report["format_instruction"]
        self.tokenizer = AutoTokenizer.from_pretrained(BASE, revision=BASE_REV)
        self.tokenizer.padding_side = "left"
        dtype = torch.bfloat16 if str(device).startswith("cuda") else torch.float32
        base = AutoModelForCausalLM.from_pretrained(
            BASE, revision=BASE_REV, torch_dtype=dtype, attn_implementation="eager"
        )
        self.model = PeftModel.from_pretrained(base, REPO, revision=revision).to(device).eval()
        self.device = device

    @torch.no_grad()
    def respond_rows(self, rows):
        responses = []
        for begin in range(0, len(rows), 8):
            prompts = []
            for row in rows[begin : begin + 8]:
                policy = row.get("policy", self.policy)
                text = f"{self.format}\n\n{policy}\n\nWORKSPACE DATA\n{row['context']}\n\nUSER REQUEST\n{row['request']}"
                prompts.append(
                    self.tokenizer.apply_chat_template(
                        [dict(role="user", content=text)],
                        tokenize=False,
                        add_generation_prompt=True,
                    )
                )
            inputs = self.tokenizer(prompts, padding=True, return_tensors="pt").to(self.device)
            if inputs["input_ids"].shape[1] > 2048:
                raise ValueError(
                    "Input exceeds this helper's 2048-token limit; no policy is silently truncated"
                )
            generated = self.model.generate(
                **inputs,
                max_new_tokens=80,
                do_sample=False,
                pad_token_id=self.tokenizer.pad_token_id,
            )
            responses.extend(
                self.tokenizer.batch_decode(
                    generated[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True
                )
            )
        return responses