"""Prepare immutable model/tokenizer assets during build, never during startup."""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path

DEFAULT_MODEL = 'intfloat/multilingual-e5-small'
DEFAULT_REVISION = '614241f622f53c4eeff9890bdc4f31cfecc418b3'
DEFAULT_ROOT = Path(__file__).resolve().parents[2] / 'runtime_assets'


def asset_root() -> Path:
    return Path(os.getenv('ZP_RUNTIME_ASSETS', str(DEFAULT_ROOT)))


def manifest(root: Path | None = None) -> dict | None:
    path = (root or asset_root()) / 'manifest.json'
    if not path.exists():
        if os.getenv('ZP_ASSETS_REQUIRED') == '1':
            raise RuntimeError('Chýbajú runtime assets; zostavte image alebo spustite python -m tools.runtime.runtime_assets prepare.')
        return None
    return json.loads(path.read_text())


def configure_tokenizer() -> None:
    data = manifest()
    if data is not None:
        # Bundled immutable files take precedence over an empty user cache.
        os.environ['TIKTOKEN_CACHE_DIR'] = str(asset_root() / 'tiktoken')


def local_model(model_name: str, revision: str | None) -> str | None:
    data = manifest()
    if data and data['embedding']['model'] == model_name:
        if revision and revision != data['embedding']['revision']:
            raise RuntimeError('Požadovaná revízia modelu nie je v image; vytvorte nový image s touto revíziou.')
        return str(asset_root() / 'model')
    if data and os.getenv('ZP_ASSETS_REQUIRED') == '1':
        raise RuntimeError('Požadovaný embeddingový model nie je v image; vytvorte nový image.')
    return None


def verify(root: Path) -> dict:
    data = json.loads((root / 'manifest.json').read_text())
    for relative, digest in data['files'].items():
        path = root / relative
        if not path.resolve().is_relative_to(root.resolve()):
            raise RuntimeError('Manifest obsahuje cestu mimo runtime assets')
        if not path.is_file() or hashlib.sha256(path.read_bytes()).hexdigest() != digest:
            raise RuntimeError(f'Chýbajúci alebo poškodený runtime asset: {relative}')
    if not any(name.startswith('tiktoken/') for name in data['files']):
        raise RuntimeError('Manifest neobsahuje tokenizér')
    required = {'model/config.json', 'model/modules.json', 'model/tokenizer.json'}
    if not required.issubset(data['files']) or not any(name.endswith('.safetensors') for name in data['files']):
        raise RuntimeError('Manifest neobsahuje kompletný embeddingový model')
    return data


def offline_probe(root: Path) -> None:
    """Fresh interpreter, with network calls forbidden even on an online builder."""
    environment = {**os.environ, 'ZP_RUNTIME_ASSETS': str(root.resolve()), 'ZP_ASSETS_REQUIRED': '1',
        'HF_HUB_OFFLINE': '1', 'TRANSFORMERS_OFFLINE': '1', 'HF_HOME': str(root / 'unused-user-cache'),
        'TIKTOKEN_CACHE_DIR': str(root / 'unused-token-cache')}
    code = '''
import socket
def denied(*args, **kwargs): raise RuntimeError('Network access during offline asset probe')
socket.socket.connect = denied
socket.socket.connect_ex = denied
socket.getaddrinfo = denied
from tools.runtime.runtime_assets import configure_tokenizer, manifest, local_model
configure_tokenizer()
import tiktoken
for encoding in manifest()['encodings']:
    assert tiktoken.get_encoding(encoding).encode('Overenie štartu bez siete.')
from sentence_transformers import SentenceTransformer
data = manifest()['embedding']
model = SentenceTransformer(local_model(data['model'], data['revision']), local_files_only=True)
assert model.encode(['query: overenie bez siete'], show_progress_bar=False).shape[0] == 1
print('Offline model/tokenizer probe: OK')
'''
    subprocess.run([sys.executable, '-c', code], env=environment, check=True)


def prepare(root: Path, model: str, revision: str | None, encodings: list[str]) -> dict:
    revision = revision or (DEFAULT_REVISION if model == DEFAULT_MODEL else '')
    if not re.fullmatch('[0-9a-f]{40}', revision):
        raise ValueError('Príprava modelu vyžaduje presnú 40-znakovú Git revíziu')
    if (root / 'manifest.json').is_file():
        data = verify(root)
        if data['embedding'] == {'model': model, 'revision': revision} and data['encodings'] == encodings:
            offline_probe(root)
            return data
        raise RuntimeError('Assets patria inej konfigurácii; použite nový/prázdny pracovný adresár.')
    from huggingface_hub import snapshot_download
    from huggingface_hub.errors import LocalEntryNotFoundError
    try:
        source = snapshot_download(model, revision=revision, local_files_only=True)
    except (FileNotFoundError, LocalEntryNotFoundError):
        source = snapshot_download(model, revision=revision,
            allow_patterns=['*.json', '*.safetensors', '*.model', '*.txt'],
            ignore_patterns=['onnx/*', 'openvino/*'])
    root.mkdir(parents=True, exist_ok=True)
    shutil.copytree(source, root / 'model', dirs_exist_ok=True, symlinks=False,
        ignore=shutil.ignore_patterns('onnx', 'openvino', '*.bin', '*.md', '.cache', '.gitattributes'))
    environment = {**os.environ, 'TIKTOKEN_CACHE_DIR': str(root.resolve() / 'tiktoken')}
    subprocess.run([sys.executable, '-c',
        'import sys,tiktoken; [tiktoken.get_encoding(name) for name in sys.argv[1:]]', *encodings],
        env=environment, check=True)
    files = {str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest()
             for folder in ['model', 'tiktoken'] for p in sorted((root / folder).rglob('*')) if p.is_file()}
    data = {'schema_version': 1, 'embedding': {'model': model, 'revision': revision},
            'encodings': encodings, 'files': files}
    (root / 'manifest.json').write_text(json.dumps(data, sort_keys=True, indent=2) + '\n')
    verify(root)
    offline_probe(root)
    return data


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('command', choices=['prepare', 'verify'])
    parser.add_argument('--output', type=Path, default=asset_root())
    parser.add_argument('--model', default=DEFAULT_MODEL)
    parser.add_argument('--revision')
    parser.add_argument('--encoding', action='append')
    args = parser.parse_args()
    if args.command == 'prepare':
        data = prepare(args.output, args.model, args.revision, list(dict.fromkeys(args.encoding or ['cl100k_base'])))
    else:
        data = verify(args.output)
        offline_probe(args.output)
    print(json.dumps({'model': data['embedding'], 'encodings': data['encodings'], 'files': len(data['files'])}))


if __name__ == '__main__':
    main()