"""Hugging Face Inference Endpoints custom handler for Banyan Sapling.

Loads the model exactly like the reference `inference.py` shipped with the
weights package (verified against the H200 reference runtime), with two
serving-only deviations:
  - skips the strict package-version assertion loop (requirements.txt pins them)
  - skips the full 46 GB sha256 pre-hash (the Hub delivers checksummed storage)
"""
from __future__ import annotations

import json
from pathlib import Path
import sys
import time
from typing import Any, Dict

CURRENT_DIR = Path(__file__).resolve().parent
if str(CURRENT_DIR) not in sys.path:
    sys.path.insert(0, str(CURRENT_DIR))

import os
os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8')


def load_model_for_endpoint(root: Path, device: str = "cuda"):
    root = Path(root).resolve()
    if str(root) not in sys.path:
        sys.path.insert(0, str(root))

    import torch
    from safetensors.torch import load_file
    # The text-only runtime does not require torchvision or torchaudio.
    sys.modules.setdefault('torchvision', None)
    sys.modules.setdefault('torchaudio', None)
    from transformers import AutoConfig, AutoModel, AutoTokenizer
    from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeRotaryEmbedding
    from runtime.backbone import attach_lora, set_shared
    from runtime.model import FullModel, identity, seed_all, tensor_hash
    from runtime.conformal_calibrator import ConformalCalibrator

    seed_all(1707)
    config_data = json.loads((root / 'config.json').read_text())
    architecture = json.loads((root / 'architecture.json').read_text())
    expert_counts = [config_data[k] for k in ('num_local_experts', 'num_experts') if k in config_data]
    if (not expert_counts or any(n != 96 for n in expert_counts)
            or config_data.get('model_type') != 'qwen3_moe'
            or identity(config_data) != architecture['hf_config_sha256']
            or architecture.get('modify_rope') is not False
            or architecture.get('factorized_lora') is not True):
        raise ValueError('Unsupported model architecture')
    config = AutoConfig.from_pretrained(root, local_files_only=True)
    with torch.device('meta'):
        backbone = AutoModel.from_config(config, dtype=torch.bfloat16,
                                        attn_implementation='sdpa', experts_implementation='grouped_mm')
        attach_lora(backbone, architecture['config'])
        if architecture['shared_ids'] is not None:
            set_shared(backbone, architecture['shared_ids'])
        model = FullModel(backbone)
    state = load_file(str(root / 'model.safetensors'), device=device)
    model.load_state_dict(state, strict=True, assign=True)
    del state
    model.backbone.rotary_emb = Qwen3MoeRotaryEmbedding(config, device=device)
    rope = model.backbone.rotary_emb
    scaling = rope.attention_scaling
    fingerprint = {'inv_freq': tensor_hash({'inv_freq': rope.inv_freq}),
                   'attention_scaling': scaling.item() if hasattr(scaling, 'item') else scaling}
    if fingerprint != architecture['rotary_fingerprint']:
        raise ValueError('Native positional encoding does not match the checkpoint')
    if any(t.is_meta for t in list(model.parameters()) + list(model.buffers())):
        raise ValueError('Incomplete model materialization')
    for name, parameter in model.named_parameters():
        expected = torch.float32 if name.startswith('head.') or name.endswith(('.A', '.B')) or '.mlp.gate.' in name else torch.bfloat16
        if parameter.dtype != expected:
            raise ValueError(f'Unexpected precision: {name}')
    model.requires_grad_(False)
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(root / 'tokenizer', local_files_only=True)
    calibrator = ConformalCalibrator.from_dict(json.loads((root / 'calibration.json').read_text()))

    manifest_path = root / 'distribution_manifest.json'
    if manifest_path.exists():
        manifest = json.loads(manifest_path.read_text())
    else:
        manifest = {'source_release_sha256': '', 'files': {'model.safetensors': ''}}
    return model, tokenizer, calibrator, manifest


def predict_for_endpoint(model, tokenizer, calibrator, manifest, request):
    import torch
    from runtime.backbone import encode
    from runtime.model import infer

    row = encode(tokenizer, request, 16384)
    pad = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
    start = time.perf_counter()
    logits = infer(model, [row], pad)[0]
    detail = calibrator.predict_detailed(logits.reshape(1, 4), empty_set_policy='rounded_argmax')[0]
    probabilities = torch.softmax(logits.double() / calibrator.temperature, -1).tolist()
    return {'model': 'banyan-sapling', 'release_stage': 'research-preview',
            'source_release_sha256': manifest.get('source_release_sha256', ''),
            'weights_sha256': manifest.get('files', {}).get('model.safetensors', ''),
            'logits': logits.tolist(), **detail,
            'probabilities': dict(zip(calibrator.LABELS, probabilities)),
            'usage': {'input_tokens': len(row['input_ids']), 'generated_tokens': 0},
            'timing_ms': {'model': (time.perf_counter() - start) * 1000}}


class EndpointHandler:
    def __init__(self, path: str = "") -> None:
        model_dir = Path(path) if path else CURRENT_DIR
        self.model, self.tokenizer, self.calibrator, self.manifest = load_model_for_endpoint(
            model_dir, device="cuda"
        )

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        payload = data.get("inputs", data)
        if isinstance(payload, str):
            try:
                payload = json.loads(payload)
            except Exception:
                return {"error": "Invalid payload: expected JSON object with policy, evidence, and claim strings."}

        if not isinstance(payload, dict):
            return {"error": "Payload must be a dictionary with policy, evidence, and claim fields."}

        for field in ("policy", "evidence", "claim"):
            val = payload.get(field)
            if not val or not str(val).strip():
                return {"error": f"Missing required non-empty string field: '{field}'"}

        return predict_for_endpoint(self.model, self.tokenizer, self.calibrator, self.manifest, payload)