import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import app

WEIGHTS = {}
MODELS = {}
for _name, _spec in app.MODEL_REGISTRY.items():
    WEIGHTS[_spec['file']] = _spec
    MODELS[_name] = _spec
    if _spec.get('dni'):
        WEIGHTS[_spec['dni']['file']] = _spec['dni']
for _name, _spec in app.FACE_REGISTRY.items():
    WEIGHTS[_spec['file']] = _spec
    MODELS[_name] = _spec
for _name, _spec in app.FACEXLIB_WEIGHTS.items():
    WEIGHTS[_name] = _spec


def wanted():
    raw = os.environ.get('PHOTO_AI_PREFETCH', '').strip()
    if not raw or raw.lower() == 'all':
        return dict(WEIGHTS)
    picked = {}
    for key in raw.split(','):
        name = key.strip()
        if not name:
            continue
        if name.lower() == 'face':
            for face_name, face_spec in app.FACE_REGISTRY.items():
                picked[face_spec['file']] = face_spec
            continue
        spec = MODELS.get(name) or WEIGHTS.get(name)
        if spec is None:
            sys.stderr.write('неизвестная модель для предзагрузки: %s\n' % name)
            continue
        picked[spec['file']] = spec
    return picked


def fetch(name, spec, target, cache_dir):
    url = spec['url']
    if cache_dir:
        cached = os.path.join(cache_dir, name)
        if app.file_ok(cached, spec['min_bytes']):
            app.link_or_copy(cached, target)
            print('веса %s взяты из кэша сборки %s' % (name, cached))
            return
        try:
            app.download_weight(url, cached, spec['min_bytes'])
        except Exception:
            if os.path.isfile(cached):
                os.unlink(cached)
            raise
        app.link_or_copy(cached, target)
        return
    app.download_weight(url, target, spec['min_bytes'])


def main():
    target_dir = sys.argv[1] if len(sys.argv) > 1 else app.SEED_DIR
    cache_dir = sys.argv[2] if len(sys.argv) > 2 else ''
    if cache_dir:
        os.makedirs(cache_dir, exist_ok=True)
    os.makedirs(target_dir, exist_ok=True)
    failed = []
    for name, spec in wanted().items():
        target = os.path.join(target_dir, name)
        if app.file_ok(target, spec['min_bytes']):
            continue
        try:
            fetch(name, spec, target, cache_dir)
        except Exception as err:
            failed.append('%s: %s' % (name, err))
            continue
    for line in failed:
        sys.stderr.write('не удалось скачать %s\n' % line)
    return 1 if failed else 0


if __name__ == '__main__':
    sys.exit(main())