"""Standalone Nima inference. Imported MLX and weights only inside load/predict."""
import argparse
import hashlib
import json
from pathlib import Path
import time
from encoding import RULES, choices, encode
from head import create_head, pack

HERE = Path(__file__).resolve().parent


def digest(path):
    h = hashlib.sha256()
    with Path(path).open('rb') as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b''):
            h.update(block)
    return h.hexdigest()


def load(release_dir=HERE, local_files_only=False):
    import mlx.core as mx
    import mlx.nn as nn
    from mlx.utils import tree_flatten
    from mlx_lm import load as load_qwen
    from huggingface_hub import snapshot_download

    folder = Path(release_dir)
    config = json.loads((folder / 'config.json').read_text())
    architecture = json.loads((folder / 'architecture.json').read_text())
    head_path = folder / config['head_file']
    if digest(head_path) != config['head_sha256']:
        raise ValueError('Nima head checksum mismatch')
    snapshot = Path(snapshot_download(
        config['base_model'], revision=config['base_revision'],
        allow_patterns=list(config['base_file_sha256']),
        local_files_only=local_files_only,
    ))
    for name, expected in config['base_file_sha256'].items():
        if digest(snapshot / name) != expected:
            raise ValueError(f'Frozen backbone file checksum mismatch: {name}')
    model, tokenizer, base_config = load_qwen(str(snapshot), return_config=True)
    if model.model_type != 'qwen3' or base_config.get('hidden_size') != architecture['hidden_size']:
        raise ValueError('Expected the pinned Qwen3 architecture')
    if base_config.get('quantization', {}).get('bits') != 4:
        raise ValueError('Expected a 4-bit backbone')
    if not any(isinstance(module, nn.QuantizedLinear) for _, module in model.named_modules()):
        raise ValueError('Quantized backbone layers missing')
    model.freeze()
    model.eval()
    if tree_flatten(model.trainable_parameters()):
        raise ValueError('Backbone must be completely frozen')
    head = create_head(architecture)
    head.load_weights(str(head_path))
    head.eval()
    if sum(value.size for _, value in tree_flatten(head.parameters())) != config['head_parameters']:
        raise ValueError('Head parameter count mismatch')
    mx.eval(model.parameters(), head.parameters())
    return Nima(model, tokenizer, head, architecture, config)


class Nima:
    def __init__(self, model, tokenizer, head, architecture, config):
        self.model, self.tokenizer, self.head = model, tokenizer, head
        self.architecture, self.config = architecture, config

    def predict(self, request):
        import mlx.core as mx
        if not isinstance(request.get('prompt'), str) or not request['prompt'].strip():
            raise ValueError('A nonempty prompt is required')
        if not isinstance(request.get('context'), dict) or not {'current_date', 'timezone'} <= request['context'].keys():
            raise ValueError('context requires current_date and timezone')
        names = choices(request)
        if 'no_tool' not in names:
            raise ValueError('Include an explicit no_tool choice to allow abstention')
        case = {**request, 'id': 'custom', 'target': names[0]}
        started = time.perf_counter()
        encoded = encode(case, names, self.tokenizer, RULES, self.architecture['max_tokens'])
        # Torso only: Model.__call__, vocabulary projection and generation are never used.
        hidden = self.model.model(mx.array([encoded['token_ids']], dtype=mx.int32), cache=None)
        evidence = mx.stop_gradient(hidden[0, mx.array(encoded['evidence_positions']), :].astype(mx.float32))
        candidates = mx.stack([mx.mean(hidden[0, span['start']:span['end'], :], axis=0) for span in encoded['candidate_spans']])
        candidates = mx.stop_gradient(candidates.astype(mx.float32))
        mx.eval(evidence, candidates)
        q, e, qv, ev, _ = pack([{'candidates': candidates, 'evidence': evidence, 'label': 0}])
        logits = self.head(q, e, qv, ev)[0, :len(names)]
        probabilities = mx.softmax(logits / self.config['temperature'])
        mx.eval(logits, probabilities)
        elapsed = time.perf_counter() - started
        scores = dict(zip(names, probabilities.tolist()))
        winner = min(scores, key=lambda name: (-scores[name], name))
        return {'prediction': winner, 'probabilities': scores,
                'logits': dict(zip(names, logits.tolist())), 'latency_seconds': elapsed,
                'input_tokens': len(encoded['token_ids']), 'qwen_forward_passes': 1,
                'note': 'Probabilities are uncalibrated and relative to the supplied choices; no tool executes.'}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--input', type=Path, required=True)
    parser.add_argument('--local-files-only', action='store_true')
    args = parser.parse_args()
    request = json.loads(args.input.read_text())
    router = load(local_files_only=args.local_files_only)
    print(json.dumps(router.predict(request), indent=2, ensure_ascii=False))


if __name__ == '__main__':
    main()