"""Shared local loaders and provenance helpers; no downloads or publication."""
from __future__ import annotations

import hashlib
import importlib.util
from importlib.metadata import distribution, version, PackageNotFoundError
import json
from pathlib import Path
import random
import sys


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


def model_fingerprint(model_dir, backend):
    model_dir = Path(model_dir)
    names = {
        'torch': ('model.safetensors', 'encoder/config.json', 'tokenizer/tokenizer.json', 'rl_agent_config.json'),
        'onnx': ('laya.onnx', 'laya.onnx.data', 'tokenizer.json', 'rl_agent_config.json'),
        'mlx': ('mlx/model.safetensors', 'mlx/config.json', 'tokenizer.json', 'rl_agent_config.json'),
    }[backend]
    return {name: sha256(model_dir / name) for name in names}


def versions():
    result = {'python': sys.version.split()[0]}
    for name in ('torch', 'transformers', 'laya', 'mlx', 'onnxruntime', 'onnx',
                 'onnxscript', 'tokenizers', 'safetensors', 'numpy'):
        try:
            result[name] = version(name)
        except PackageNotFoundError:
            continue
    try:
        direct = distribution('laya').read_text('direct_url.json')
        if direct:
            result['laya_commit'] = json.loads(direct).get('vcs_info', {}).get('commit_id')
    except PackageNotFoundError:
        pass
    return result


def require_pinned_reference():
    installed = versions()
    expected = {'torch': '2.14.0', 'transformers': '5.17.0', 'laya': '0.3.20'}
    wrong = {name: installed.get(name) for name, value in expected.items()
             if installed.get(name, '').split('+')[0] != value}
    commit = installed.get('laya_commit')
    if commit and commit != '4066d5d5fbf08b66c6757ddeedbd797bd7655bc0':
        wrong['laya_commit'] = commit
    if wrong:
        raise RuntimeError(f'Reference environment differs from the trained checkpoint: {wrong}. '
                           'Use Python 3.12 and pip install -r requirements.txt.')


def load_module(path, name):
    spec = importlib.util.spec_from_file_location(name, str(path))
    if spec is None or spec.loader is None:
        raise ImportError(f'Cannot load {path}')
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


class TorchSession:
    """Same session interface as ONNX/MLX, with a strict local checkpoint load."""
    def __init__(self, checkpoint, device='cpu', threads=4):
        require_pinned_reference()
        import torch
        from safetensors.torch import load_file
        from laya.common import build_model

        checkpoint = Path(checkpoint)
        self.torch = torch
        self.device = torch.device(device)
        torch.set_num_threads(threads)
        cfg = json.loads((checkpoint / 'rl_agent_config.json').read_text())
        self.model = build_model(cfg, encoder_dir=str(checkpoint / 'encoder'))
        self.model.encoder.config.reference_compile = False
        self.model.load_state_dict(load_file(str(checkpoint / 'model.safetensors')), strict=True)
        self.model.float().eval().to(self.device)

    def run(self, names, feeds):
        if list(names) != ['logits']:
            raise ValueError(f'Unsupported outputs: {names}')
        torch = self.torch
        with torch.inference_mode():
            f = {key: torch.as_tensor(value, device=self.device,
                                      dtype=torch.bool if key == 'marker_mask' else torch.long)
                 for key, value in feeds.items()}
            return [self.model(**f)[0].float().cpu().numpy()]


def load_router(model_dir, backend='onnx', repo_dir=None, threads=4, device='cpu'):
    """Load checkpoint or deployment using identical production prompt/decision logic."""
    model_dir = Path(model_dir).resolve()
    module_dir = Path(__file__).resolve().parent
    project_root = module_dir.parent if module_dir.name == 'scripts' else module_dir
    repo_dir = Path(repo_dir) if repo_dir else project_root / 'laya-intent-router'
    candidates = (repo_dir / 'laya_intent_router.py', Path(__file__).parent / 'laya_intent_router.py',
                  model_dir / 'laya_intent_router.py')
    runtime = next((path for path in candidates if path.is_file()), None)
    if runtime is None:
        raise FileNotFoundError('laya_intent_router.py is missing; specify repo_dir or use a complete release package')
    mod = load_module(runtime, 'lapis_router_runtime')
    if backend != 'torch':
        return mod.LayaIntentRouter(str(model_dir), backend=backend, threads=threads, shortlist_k=0)

    from tokenizers import Tokenizer
    router = mod.LayaIntentRouter.__new__(mod.LayaIntentRouter)
    router.cfg = json.loads((model_dir / 'rl_agent_config.json').read_text())
    router.tok = Tokenizer.from_file(str(model_dir / 'tokenizer/tokenizer.json'))
    router.cls, router.sep, router.mask = (router.tok.token_to_id(t) for t in ('[CLS]', '[SEP]', '[MASK]'))
    router.backend = 'torch'
    router.sess = TorchSession(model_dir, device=device, threads=threads)
    router.threshold = router.cfg.get('match_threshold', 0.625)
    router.temps = router.cfg['temperature_by_options']
    router.k, router.shortlist = 0, None
    return router


def load_episode_samples(path, n, seed):
    if n < 1:
        raise ValueError('Sample count must be positive')
    eps = [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()]
    if not eps:
        raise ValueError(f'No episodes in {path}')
    if any(str(ep.get('split', '')).lower() == 'test' for ep in eps):
        raise ValueError('Use development episodes, not the reserved test split')
    return random.Random(seed).sample(eps, min(n, len(eps)))


def episode_paths(ep):
    paths = ep['paths']
    if isinstance(paths, dict):
        return {key: [value] if isinstance(value, str) else value for key, value in paths.items()}
    result = {}
    for path in paths:
        key = path.get('path_id') or path.get('id') or path.get('name')
        value = path.get('text') or path.get('texts') or path.get('descriptions')
        result[key] = [value] if isinstance(value, str) else list(value)
    return result