#!/usr/bin/env python3
"""Attach a downloaded ego-exo model bundle to code/core. Dry-run by default.

Creates links only where destinations are absent. Loader patches are applied
only to the exact known pre-sharding code version (or treated as already applied
if hashes match); other local edits are never overwritten. No inference/upload.
"""
import argparse
import hashlib
import json
from pathlib import Path
import tempfile


def digest(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def safe_relative(name):
    path = Path(name)
    if path.is_absolute() or not path.parts or '..' in path.parts or '.' in path.parts:
        raise ValueError('unsafe relative path in loader patch')
    return path


def plan_install(bundle, code_root, patch_loaders=True):
    bundle, code_root = Path(bundle).resolve(), Path(code_root).resolve()
    if not (code_root / 'project_paths.py').is_file():
        raise ValueError('--code-root must be the existing code/core directory')
    if not (bundle / 'bundle_manifest.json').is_file():
        raise ValueError('bundle is incomplete or has no verified bundle_manifest.json')
    if json.loads((bundle / 'bundle_manifest.json').read_text()).get('status') != 'complete':
        raise ValueError('bundle is not complete')
    operations = []
    for group in ('models', 'checkpoints'):
        source_root = bundle / group
        if not source_root.is_dir():
            raise ValueError('missing bundle resource group: ' + group)
        target_root = code_root / group
        if target_root.exists() and not target_root.is_dir():
            raise ValueError('resource parent is not a directory: ' + str(target_root))
        for source in sorted(source_root.iterdir()):
            if not source.is_dir():
                continue
            target = target_root / source.name
            if target.is_symlink() and target.resolve() == source.resolve():
                continue
            if target.exists() or target.is_symlink():
                raise ValueError('existing resource will not be overwritten: ' + str(target))
            operations.append(dict(action='link', source=str(source), target=str(target)))
    if patch_loaders:
        patch_root = bundle / 'code_support'
        manifest = json.loads((patch_root / 'loader_patch.json').read_text())
        for row in manifest['files']:
            relative = safe_relative(row['path'])
            source, target = patch_root / relative, code_root / relative
            if digest(source) != row['new_sha256']:
                raise ValueError('bundle loader patch hash differs')
            if target.is_symlink():
                raise ValueError('refusing to patch a symlink: ' + str(target))
            if target.exists() and not target.is_file():
                raise ValueError('refusing to patch a non-file: ' + str(target))
            current = digest(target) if target.is_file() else None
            if current == row['new_sha256']:
                continue
            if current != row['previous_sha256']:
                raise ValueError('unknown/local code changes will not be overwritten: ' + str(target))
            operations.append(dict(action='patch', source=str(source), target=str(target)))
    return operations


def install(bundle, code_root, apply=False, patch_loaders=True):
    operations = plan_install(bundle, code_root, patch_loaders)
    for row in operations:
        print(row['action'].upper(), row['target'], '<-', row['source'])
    if not apply:
        print('DRY RUN: nothing changed. Add --apply after checking the paths.')
        return operations
    # Validate the entire plan before any writes. Never replace a resource directory.
    for row in operations:
        source, target = Path(row['source']), Path(row['target'])
        target.parent.mkdir(parents=True, exist_ok=True)
        if row['action'] == 'link':
            target.symlink_to(source, target_is_directory=True)
        else:
            with tempfile.NamedTemporaryFile(dir=target.parent, prefix='.loader_patch-', delete=False) as handle:
                temporary = Path(handle.name)
                handle.write(source.read_bytes())
            if target.exists():
                temporary.chmod(target.stat().st_mode & 0o777)
            else:
                temporary.chmod(0o644)
            temporary.replace(target)
    print('ATTACHED. No training, inference, GPU use or upload occurred.')
    return operations


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--bundle', type=Path, default=Path(__file__).resolve().parents[1])
    parser.add_argument('--code-root', type=Path, required=True)
    parser.add_argument('--apply', action='store_true')
    parser.add_argument('--no-code-patch', action='store_true')
    args = parser.parse_args()
    install(args.bundle, args.code_root, args.apply, not args.no_code_patch)


if __name__ == '__main__':
    main()