import importlib.util
import json
import tempfile
import unittest
from pathlib import Path


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


def load_worker():
    specification = importlib.util.spec_from_file_location(
        "snapform_download_worker",
        PROJECT_ROOT / "src" / "snapform" / "worker.py",
    )
    module = importlib.util.module_from_spec(specification)
    specification.loader.exec_module(module)
    return module


class FakeTqdm:
    def __init__(self, *args, **kwargs):
        self.n = kwargs.get("initial", 0)
        self.total = kwargs.get("total", 0)

    def update(self, amount=1):
        self.n += amount
        return True

    def refresh(self, *args, **kwargs):
        return True


class DownloadProgressTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.worker = load_worker()

    def test_reports_aggregated_download_bytes(self):
        with tempfile.TemporaryDirectory() as directory:
            status_path = Path(directory) / "status.json"
            progress_class = self.worker.create_download_progress_class(
                status_path,
                "owner/model",
                tqdm_type=FakeTqdm,
                clock=lambda: 1.0,
            )
            progress = progress_class(
                desc="Downloading bytes",
                total=100,
                initial=0,
                unit="B",
            )

            progress.update(25)

            status = json.loads(status_path.read_text(encoding="utf-8"))
            self.assertAlmostEqual(status["progress"], 0.275)
            self.assertEqual(status["message"], "Downloading owner/model")

    def test_ignores_unrelated_progress_bars(self):
        with tempfile.TemporaryDirectory() as directory:
            status_path = Path(directory) / "status.json"
            progress_class = self.worker.create_download_progress_class(
                status_path,
                "owner/model",
                tqdm_type=FakeTqdm,
                clock=lambda: 1.0,
            )
            progress = progress_class(
                desc="Reconstructing",
                total=100,
                initial=0,
                unit="B",
            )

            progress.update(25)

            self.assertFalse(status_path.exists())


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