"""Provision an offline, verified EmbeddingGemma 2 Windows distribution."""
import argparse
import hashlib
import json
from pathlib import Path
import shutil
import tempfile

ROOT = Path(__file__).resolve().parents[1]
DEST = ROOT / '.runtime/windows/semantic'
PIN = 'bfcd298762cc34d0357ece5ebdd31791a3a374d8'
MODELS = {
    'embeddinggemma-2-Q8_0.gguf': (309855456, '2188ac1deca4b77dffefd603c2776a9d76d9d74ec01841392982ebb840b09135'),
    'mmproj-embeddinggemma-2-Q8_0.gguf': (554821024, 'c4a8a52691ecef40618438928bdf9e68379b854e24166f292592353db0aab64f'),
}

def verify(base, item):
    name = item['path']
    if Path(name).name != name or ':' in name or '\\' in name or '/' in name:
        raise ValueError('Invalid artifact name')
    file = base / name
    if file.is_symlink() or not file.is_file() or file.stat().st_size != item['bytes']:
        raise ValueError('Artifact missing or incorrect size: ' + name)
    digest = hashlib.sha256()
    with file.open('rb') as stream:
        while block := stream.read(1024 * 1024): digest.update(block)
    if digest.hexdigest() != item['sha256']: raise ValueError('Artifact digest mismatch: ' + name)

def validate(base):
    manifest = json.loads((base / 'models/model-manifest.json').read_text('utf-8'))
    if (manifest['model_id'], manifest['revision'], manifest['dimensions'], manifest['quantization'], manifest['runtime_tag']) != ('google/embeddinggemma-2', PIN, 768, 'Q8_0', 'b11461'):
        raise ValueError('Model identity mismatch')
    for name, (size, digest) in MODELS.items(): verify(base / 'models', {'path': name, 'bytes': size, 'sha256': digest})
    files = json.loads((base / 'runtime-manifest.json').read_text('utf-8'))['files']
    if not 1 <= len(files) <= 64 or len({item['path'].lower() for item in files}) != len(files): raise ValueError('Invalid runtime inventory')
    if not any(item['path'] == 'searchi-gemma-worker.exe' for item in files): raise ValueError('Worker absent')
    for item in files: verify(base, item)
    if {p.name.lower() for p in base.glob('*.dll')} != {i['path'].lower() for i in files if i['path'].lower().endswith('.dll')}: raise ValueError('Unverified DLL')

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--source-project', type=Path)
    parser.add_argument('--verify', action='store_true')
    args = parser.parse_args()
    if DEST.exists():
        validate(DEST)
        print('Verified offline EmbeddingGemma 2 model and Windows runtime.')
        return
    if args.verify or args.source_project is None: raise ValueError('Provision semantic runtime with --source-project before packaging')
    source = args.source_project.resolve()
    runtime = source / '.searchi/runtime/auto/Release'
    inventory = json.loads((runtime / 'runtime-libraries.json').read_text('utf-8'))
    worker = runtime / 'searchi-gemma-worker.exe'
    files = [{'path': worker.name, 'bytes': worker.stat().st_size, 'sha256': inventory['worker_sha256']}, *inventory['libraries']]
    for item in files: verify(runtime, item)
    DEST.parent.mkdir(parents=True, exist_ok=True)
    candidate = Path(tempfile.mkdtemp(prefix='semantic-stage-', dir=DEST.parent))
    # A failed stage is retained for diagnosis; never delete another project's assets.
    (candidate / 'models').mkdir()
    for item in files: shutil.copy2(runtime / item['path'], candidate / item['path'])
    (candidate / 'runtime-manifest.json').write_text(json.dumps({'files': files}, indent=2), 'utf-8')
    shutil.copy2(source / 'tools/embeddinggemma2/model-manifest.json', candidate / 'models/model-manifest.json')
    for name in MODELS: shutil.copy2(source / '.searchi/models/embeddinggemma2' / name, candidate / 'models' / name)
    shutil.copy2(source / '.searchi/references/llama.cpp/LICENSE', candidate / 'LLAMA-LICENSE.txt')
    for name in ('MSVC-REDISTRIBUTION.txt', 'VULKAN-LICENSE.txt'): shutil.copy2(runtime / name, candidate / name)
    validate(candidate)
    candidate.rename(DEST)
    print('Provisioned verified offline EmbeddingGemma 2 model and Windows runtime.')

if __name__ == '__main__': main()