import importlib.util
import sys
import unittest
from pathlib import Path
from unittest.mock import patch


PROJECT_ROOT = Path(__file__).resolve().parents[1]


def load_sync_models():
    path = PROJECT_ROOT / "tools" / "models" / "cli.py"
    spec = importlib.util.spec_from_file_location("sync_models", path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


class SyncModelsTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.sync = load_sync_models()

    def test_r2_is_only_an_upload_destination(self):
        for model in self.sync.MODEL_FILES:
            self.assertNotIn("models.omoolab.xyz", model.source_url)
            if model.build:
                self.assertFalse(model.source_url)

    def test_local_exports_use_official_pytorch_weights(self):
        builds = self.sync.MODEL_BUILDS
        self.assertEqual(
            set(builds),
            {"REALESRGAN_GENERAL_WDN_X4V3", "SAFMN_REAL_X4"},
        )
        for build in builds.values():
            self.assertTrue(build.source_url.startswith("https://github.com/"))
            self.assertTrue(build.source_filename.endswith(".pth"))
            self.assertEqual(len(build.source_sha256), 64)

    def test_prepare_model_builds_missing_local_export(self):
        model = next(model for model in self.sync.MODEL_FILES if model.build)
        with (
            patch.object(self.sync, "verify_file", return_value=False),
            patch.object(self.sync, "build_model") as build,
            patch.object(self.sync, "download_file") as download,
        ):
            self.sync.prepare_model(model)
        build.assert_called_once_with(model)
        download.assert_not_called()

    def test_prepare_model_downloads_upstream_onnx(self):
        model = next(
            model
            for model in self.sync.MODEL_FILES
            if model.source_url and not model.build
        )
        with (
            patch.object(self.sync, "verify_file", return_value=False),
            patch.object(self.sync, "download_file") as download,
        ):
            self.sync.prepare_model(model)
        download.assert_called_once_with(
            model.source_url,
            model.destination,
            model.size,
            model.sha256,
        )

    def test_upload_copies_declared_files_without_remote_deletion(self):
        model = self.sync.ModelFile(
            key="TEST",
            folder="folder",
            filename="model.onnx",
            size=1,
            sha256="hash",
        )
        with (
            patch.object(self.sync.shutil, "which", return_value="rclone"),
            patch.object(self.sync.subprocess, "run") as run,
        ):
            self.sync.upload_models((model,))
        command = run.call_args.args[0]
        self.assertEqual(command[1], "copyto")
        self.assertNotIn("sync", command)
        self.assertEqual(command[-2], model.remote)

    def test_verify_file_checks_size_and_sha256(self):
        path = PROJECT_ROOT / "tools" / "models" / "cli.py"
        size = path.stat().st_size
        digest = self.sync.file_sha256(path)
        self.assertTrue(self.sync.verify_file(path, size, digest))
        with self.assertRaisesRegex(RuntimeError, "unexpected size"):
            self.sync.verify_file(path, size - 1, digest)


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