"""Inference-only formula excerpts from the source-bound canonical runtime."""
from dataclasses import dataclass
import hashlib
import json
import re
from pathlib import Path
import torch
from torch import nn
MAX_TOKENS = 32768
TRAIN_TOKENS = 768
def file_sha256(path):
    digest = hashlib.sha256()
    with Path(path).open('rb') as stream:
        for chunk in iter(lambda: stream.read(4 * 1024 * 1024), b''):
            digest.update(chunk)
    return digest.hexdigest()


def audit_loading_info(info, source_keys):
    if not isinstance(info, dict) or not all(k in info for k in ('missing_keys', 'unexpected_keys', 'mismatched_keys')):
        raise ValueError('missing loading audit fields')
    unexpected = set(info['unexpected_keys'])
    allowed = {'lm_head.weight'} if 'lm_head.weight' in source_keys else set()
    if info['missing_keys'] or info['mismatched_keys'] or info.get('error_msgs') or unexpected - allowed:
        raise ValueError(f'base loading audit failed: {info}')
    return {'missing_keys': [], 'mismatched_keys': [], 'unexpected_keys': sorted(unexpected),
            'allowed_unexpected_keys': sorted(unexpected),
            'exception_reason': ('LM vocabulary head is deliberately excluded by AutoModel; source key confirmed'
                                 if unexpected else None)}


def load_base(lock):
    """Only this boundary touches pretrained assets; local and fail-closed."""
    from safetensors import safe_open
    from transformers import AutoModel, AutoTokenizer
    if lock.get('model') != 'ibm-granite/granite-4.0-350m':
        raise ValueError('wrong locked model')
    if not isinstance(lock.get('revision'), str) or not re.fullmatch('[0-9a-f]{40}', lock['revision']):
        raise ValueError('invalid locked revision')
    sha = lock.get('model_safetensors_sha256')
    if not isinstance(sha, str) or not re.fullmatch('[0-9a-f]{64}', sha):
        raise ValueError('invalid source hash')
    local = Path(lock['local_path'])
    if not local.is_absolute() or lock.get('max_context_tokens') != MAX_TOKENS:
        raise ValueError('invalid local path or context lock')
    source = local / 'model.safetensors'
    actual = file_sha256(source)
    if actual != sha:
        raise ValueError('source model hash mismatch')
    with safe_open(str(source), framework='pt', device='cpu') as handle:
        source_keys = set(handle.keys())
    model, info = AutoModel.from_pretrained(str(local), local_files_only=True,
                                           trust_remote_code=False, attn_implementation='sdpa',
                                           dtype=torch.float32, output_loading_info=True)
    audit = audit_loading_info(info, source_keys)
    if model.config.model_type != 'granitemoehybrid' or model.config.hidden_size != 1024:
        raise ValueError('unexpected base architecture (requires GraniteMoeHybrid hidden=1024)')
    model.to(device='cpu', dtype=torch.float32).eval()
    tokenizer = AutoTokenizer.from_pretrained(str(local), local_files_only=True, trust_remote_code=False)
    assets = ('config.json', 'tokenizer.json', 'tokenizer_config.json',
              'special_tokens_map.json', 'vocab.json', 'merges.txt', 'added_tokens.json', 'tokenizer.model')
    asset_hashes = {name: file_sha256(local/name) for name in assets if (local/name).is_file()}
    return model, tokenizer, dict(audit, source_sha256=actual, source_file=str(source), asset_sha256=asset_hashes)


def architecture(model):
    return {'schema_version': 1, 'backbone_class': type(model.backbone).__name__,
            'backbone_config': json.loads(model.backbone.config.to_json_string(use_diff=False)),
            'hidden_size': model.backbone.config.hidden_size, 'head_hidden': 128,
            'head_activation': 'GELU', 'head_outputs': 1, 'shared_option_head': True,
            'separate_none_head': True, 'rank': 8, 'alpha': 16, 'dropout': 0,
            'adapter_targets': model.adapter_targets,
            'option_readout': 'option_segment_last_token', 'none_readout': 'final_segment_last_token',
            'prompt_version': 'context_question_json_options_decision_v1',
            'tokenization': 'encode_each_segment_add_special_tokens_false',
            'max_input_tokens': MAX_TOKENS, 'train_max_tokens': TRAIN_TOKENS,
            'dtype': 'float32', 'device': 'cpu', 'attention': 'sdpa', 'use_cache': False}

def tensor_hash(items):
    digest = hashlib.sha256()
    for name, tensor in sorted(items):
        value = tensor.detach().cpu().contiguous()
        digest.update(name.encode())
        digest.update(str(value.dtype).encode())
        digest.update(str(tuple(value.shape)).encode())
        digest.update(value.reshape(-1).view(torch.uint8).numpy().tobytes())
    return digest.hexdigest()


def frozen_hash(model):
    return tensor_hash([(n, p) for n, p in model.named_parameters() if not p.requires_grad]
                       + list(model.named_buffers()))


def trainable_state(model):
    return {n: p.detach().cpu().contiguous() for n, p in model.named_parameters() if p.requires_grad}

class LoRALinear(nn.Module):
    """Frozen linear + rank-eight update; zero B preserves initial backbone."""
    def __init__(self, base, rank=8, alpha=16):
        super().__init__()
        if not isinstance(base, nn.Linear):
            raise ValueError('LoRA target must be nn.Linear')
        self.base = base
        self.base.requires_grad_(False)
        self.rank, self.alpha = rank, alpha
        self.A = nn.Parameter(torch.empty(rank, base.in_features, dtype=base.weight.dtype))
        self.B = nn.Parameter(torch.zeros(base.out_features, rank, dtype=base.weight.dtype))
        nn.init.kaiming_uniform_(self.A, a=5 ** 0.5)

    def forward(self, x):
        return self.base(x) + nn.functional.linear(nn.functional.linear(x, self.A), self.B) * (self.alpha / self.rank)


class OptionScorer(nn.Module):
    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone
        self.backbone.requires_grad_(False)
        attention = [(i, layer.self_attn) for i, layer in enumerate(backbone.layers)
                     if hasattr(layer, 'self_attn') and layer.self_attn is not None]
        if len(attention) < 4:
            raise ValueError('backbone requires at least four attention layers')
        self.adapter_targets = []
        for index, attn in attention[-4:]:
            for name in ('q_proj', 'v_proj'):
                setattr(attn, name, LoRALinear(getattr(attn, name)))
                self.adapter_targets.append(f'layers.{index}.self_attn.{name}')
        hidden = backbone.config.hidden_size
        self.option_head = nn.Sequential(nn.Linear(hidden, 128), nn.GELU(), nn.Linear(128, 1))
        self.none_head = nn.Sequential(nn.Linear(hidden, 128), nn.GELU(), nn.Linear(128, 1))
        self.train(False)

    def train(self, mode=True):
        super().train(mode)
        self.backbone.eval()
        for module in self.backbone.modules():
            if isinstance(module, LoRALinear):
                module.training = mode
        return self

    def forward(self, encoded):
        ids = torch.tensor([encoded.input_ids], dtype=torch.long, device='cpu')
        hidden = self.backbone(input_ids=ids, attention_mask=torch.ones_like(ids),
                               use_cache=False, return_dict=True).last_hidden_state[0]
        options = self.option_head(hidden[encoded.option_positions]).squeeze(-1)
        none = self.none_head(hidden[encoded.final_position]).reshape(1)
        return torch.cat((options, none))

def validate_request(request, *, labeled=False):
    if not isinstance(request, dict):
        raise ValueError('request must be an object')
    for key in ('context', 'question'):
        if not isinstance(request.get(key), str) or (key == 'question' and not request[key].strip()):
            raise ValueError(f'{key} must be a string (question nonempty)')
    options = request.get('options')
    if not isinstance(options, list) or not 2 <= len(options) <= 16:
        raise ValueError('options must contain 2..16 choices')
    ids = []
    for option in options:
        if not isinstance(option, dict):
            raise ValueError('option must be an object')
        for key in ('id', 'description'):
            if not isinstance(option.get(key), str) or not option[key].strip():
                raise ValueError(f'option {key} must be nonempty string')
        ids.append(option['id'])
    if len(set(ids)) != len(ids):
        raise ValueError('option ids must be unique')
    if 'target' in request and request['target'] is not None and request['target'] not in ids:
        raise ValueError('target must be option id or null')
    if labeled:
        for key in ('id', 'family'):
            if not isinstance(request.get(key), str) or not request[key].strip():
                raise ValueError(f'{key} must be nonempty')
        if 'target' not in request:
            raise ValueError('labeled row requires target')
    return request


@dataclass(frozen=True)
class EncodedRequest:
    input_ids: list
    option_positions: list
    final_position: int
    text: str