"""CPU-only adapter regression tests; never import torch or run inference."""
import importlib.util
from pathlib import Path
from types import SimpleNamespace
import unittest
from unittest.mock import Mock

spec = importlib.util.spec_from_file_location('worker', Path(__file__).parents[1] / 'scripts/yuegp-worker.py')
worker = importlib.util.module_from_spec(spec)
spec.loader.exec_module(worker)


class WorkerTests(unittest.TestCase):
    def test_windows_memory_fix_preserves_profile_and_precision(self):
        options = worker.memory_profile_options(1, False, 'win32', 32 * 1024 ** 3)
        self.assertFalse(options['pinnedMemory'])
        self.assertEqual(options['profile_no'], 1)
        self.assertFalse(options['quantizeTransformer'])
        self.assertFalse(options['compile'])
        self.assertNotIn('budgets', options)
        self.assertNotIn('pinnedMemory', worker.memory_profile_options(1, False, 'linux', 32 * 1024 ** 3))
        self.assertTrue(worker.memory_profile_options(3, False, 'win32', 32 * 1024 ** 3)['quantizeTransformer'])

    def test_lyrics_preserve_words_and_normalize_ui_headings(self):
        self.assertEqual(worker.normalize_lyrics('[Pre-Chorus]\nEvery word stays\n[Outro]\n'), '[prechorus]\nEvery word stays\n\n')
        self.assertEqual(worker.normalize_lyrics('[Verse 1]\nHello'), '[verse1]\nHello\n\n')
        for lyrics in ['', 'No heading', '[Verse]\nA\n[Chorus]\nB']:
            with self.assertRaises(ValueError):
                worker.normalize_lyrics(lyrics)

    def test_codec_construction_and_weight_loading_are_explicitly_cpu(self):
        active = []
        class Device:
            def __enter__(self): active.append('cpu')
            def __exit__(self, *args): active.pop()
        codec = Mock()
        def construct(**kwargs):
            self.assertEqual(active, ['cpu'])
            return codec
        def load(path, **kwargs):
            self.assertEqual(active, ['cpu'])
            self.assertEqual(kwargs['map_location'], 'cpu')
            return {'codec_model': 'weights'}
        torch = SimpleNamespace(device=lambda name: Device() if name == 'cpu' else self.fail('GPU context'), load=load)
        ns = {'SoundStream': construct, 'OmegaConf': SimpleNamespace(load=lambda _: SimpleNamespace(generator=SimpleNamespace(config={}))) }
        self.assertIs(worker.load_codec_on_cpu(ns, torch), codec)
        codec.load_state_dict.assert_called_once_with('weights')
        codec.eval.assert_called_once()
        self.assertEqual(active, [])


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