import hashlib
import json
from pathlib import Path
import sys
import tempfile
import unittest
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from attach_model_bundle import install, plan_install


class InstallerTests(unittest.TestCase):
    def setUp(self):
        self.temporary = tempfile.TemporaryDirectory()
        self.addCleanup(self.temporary.cleanup)
        root = Path(self.temporary.name)
        self.bundle, self.code = root / 'bundle', root / 'code'
        for path in (self.bundle/'models/base', self.bundle/'checkpoints/hoi', self.bundle/'code_support', self.code):
            path.mkdir(parents=True, exist_ok=True)
        (self.code/'project_paths.py').write_text('existing code root\n')
        (self.code/'loader.py').write_text('old known version\n')
        (self.bundle/'code_support/loader.py').write_text('new shard-aware version\n')
        old = hashlib.sha256((self.code/'loader.py').read_bytes()).hexdigest()
        new = hashlib.sha256((self.bundle/'code_support/loader.py').read_bytes()).hexdigest()
        (self.bundle/'code_support/loader_patch.json').write_text(json.dumps(
            {'files': [{'path': 'loader.py', 'previous_sha256': old, 'new_sha256': new}]}))
        (self.bundle/'bundle_manifest.json').write_text('{"status":"complete"}')

    def test_dry_run_makes_no_changes(self):
        operations = install(self.bundle, self.code)
        self.assertEqual(len(operations), 3)
        self.assertEqual((self.code/'loader.py').read_text(), 'old known version\n')
        self.assertFalse((self.code/'models').exists())

    def test_apply_and_idempotent(self):
        install(self.bundle, self.code, apply=True)
        self.assertEqual((self.code/'models/base').resolve(), self.bundle/'models/base')
        self.assertEqual((self.code/'checkpoints/hoi').resolve(), self.bundle/'checkpoints/hoi')
        self.assertEqual((self.code/'loader.py').read_text(), 'new shard-aware version\n')
        self.assertEqual(plan_install(self.bundle, self.code), [])

    def test_local_code_changes_rejected_before_any_write(self):
        (self.code/'loader.py').write_text('important local edits\n')
        with self.assertRaisesRegex(ValueError, 'local code changes'):
            install(self.bundle, self.code, apply=True)
        self.assertFalse((self.code/'models').exists())
        self.assertEqual((self.code/'loader.py').read_text(), 'important local edits\n')

    def test_existing_weight_directory_preserved(self):
        (self.code/'models/base').mkdir(parents=True)
        (self.code/'models/base/my_weight').write_text('keep')
        with self.assertRaisesRegex(ValueError, 'not be overwritten'):
            install(self.bundle, self.code, apply=True)
        self.assertEqual((self.code/'models/base/my_weight').read_text(), 'keep')

    def test_unsafe_patch_path_rejected(self):
        path = self.bundle/'code_support/loader_patch.json'
        value = json.loads(path.read_text())
        value['files'][0]['path'] = '../outside.py'
        path.write_text(json.dumps(value))
        with self.assertRaisesRegex(ValueError, 'unsafe'):
            plan_install(self.bundle, self.code)


if __name__ == '__main__':
    unittest.main()