"""Portable NF4 loader for xuhaodev/Qwen3.5-4B-Jev."""
import hashlib
import json
from pathlib import Path

import torch
from decision_model import GeneralDecisionModel
from decision_schema import SystemOneRequest, answer
from huggingface_hub import snapshot_download


def verify(directory, files):
    root = Path(directory).resolve()
    for name, expected in files.items():
        path = (root / name).resolve()
        if not path.is_relative_to(root):
            raise ValueError(f"Invalid artifact path: {name}")
        with path.open('rb') as stream:
            if hashlib.file_digest(stream, 'sha256').hexdigest() != expected:
                raise ValueError(f"Artifact hash mismatch: {name}")


class JevModel:
    def __init__(self, directory, *, base_path=None, device='cuda'):
        self.directory = Path(directory)
        manifest = json.loads((self.directory / 'manifest.json').read_text())
        verify(self.directory, manifest['files'])
        self.config = json.loads((self.directory / 'config.json').read_text())
        self.temperatures = json.loads(
            (self.directory / 'calibration.json').read_text())['temperatures']
        if any(not 0 < t < float('inf') for t in self.temperatures.values()):
            raise ValueError('Invalid temperatures')
        if not str(device).startswith('cuda') or not torch.cuda.is_available():
            raise ValueError('This NF4/FLA release requires a supported NVIDIA CUDA GPU')
        if base_path is None:
            base_path = snapshot_download(
                self.config['base_model'], revision=self.config['base_revision'],
                allow_patterns=['*.json', '*.safetensors', '*.jinja', '*.txt', 'LICENSE'])
        verify(base_path, manifest['base_files'])
        self.model = GeneralDecisionModel(
            base_path, adapter=self.directory, mode='baseline', device=device).eval()

    @classmethod
    def from_pretrained(cls, repo_id='xuhaodev/Qwen3.5-4B-Jev', *, revision=None, **kwargs):
        directory = (repo_id if Path(repo_id).is_dir()
                     else snapshot_download(repo_id, revision=revision))
        return cls(directory, **kwargs)

    @torch.inference_mode()
    def predict(self, state, questions):
        request = SystemOneRequest(state=state, model=self.config['model_name'],
                                   questions=questions)
        if len(request.questions) > 32:
            raise ValueError('At most 32 questions per request')
        answers, tokens, prepared = {}, 0, []
        for name, question in request.questions.items():
            q = question.model_dump(exclude_none=True)
            inputs = self.model.encode(state, q, max_length=self.config['max_length'])
            count = sum(item['input_ids'].shape[-1] for item in inputs)
            if count > self.config['max_expanded_tokens']:
                raise ValueError('Question expanded token budget exceeded')
            tokens += count
            if tokens > self.config.get('max_request_expanded_tokens', 262144):
                raise ValueError('Request expanded token budget exceeded')
            prepared.append((name, q, inputs))
        for name, q, inputs in prepared:
            logits = self.model.logits(inputs)
            p = torch.softmax(logits.float() / self.temperatures[q['type']], -1)
            answers[name] = answer(q, p.cpu().tolist())
        return {'model': self.config['model_name'], 'answers': answers,
                'usage': {'input_tokens': tokens, 'output_tokens': 0}}

    system_one = predict