"""CPU-only regression tests for W&B fallback with Cosmos's frozen config."""
import unittest
from unittest.mock import Mock, patch

import wandb
from cosmos_framework.utils.config import Config, JobConfig

from scripts.train_widowx_zeva import with_wandb_offline_fallback


def frozen_config(mode='online'):
    config = Config(
        model={}, optimizer={}, scheduler={}, dataloader_train=None,
        dataloader_val=None,
        job=JobConfig(project='Zeva', group='widowx', name='ipec_personal',
                      wandb_mode=mode),
    )
    config.freeze()
    return config


class WandbFallbackTest(unittest.TestCase):
    def test_online_success_does_not_retry(self):
        initialize = Mock(return_value='online')
        config = frozen_config()
        self.assertEqual(with_wandb_offline_fallback(initialize)(config, None), 'online')
        initialize.assert_called_once_with(config, None)

    def test_retry_preserves_identity_and_frozen_config(self):
        initialize = Mock(side_effect=[wandb.errors.CommError('network failure'), 'offline'])
        config = frozen_config()
        with patch('wandb.teardown') as teardown:
            self.assertEqual(with_wandb_offline_fallback(initialize)(config, None), 'offline')
        teardown.assert_called_once_with()
        retry = initialize.call_args.args[0]
        self.assertEqual(retry.job.wandb_mode, 'offline')
        self.assertEqual(retry.job.path_local, config.job.path_local)
        self.assertIs(retry.checkpoint, config.checkpoint)
        self.assertEqual(config.job.wandb_mode, 'online')
        self.assertTrue(config._is_frozen and config.job._is_frozen)

    def test_offline_failure_is_not_retried(self):
        initialize = Mock(side_effect=wandb.errors.CommError('offline failure'))
        with self.assertRaises(wandb.errors.CommError):
            with_wandb_offline_fallback(initialize)(frozen_config('offline'), None)
        self.assertEqual(initialize.call_count, 1)

    def test_unrelated_error_is_not_hidden(self):
        initialize = Mock(side_effect=RuntimeError('training failure'))
        with self.assertRaises(RuntimeError):
            with_wandb_offline_fallback(initialize)(frozen_config(), None)
        self.assertEqual(initialize.call_count, 1)

    def test_failed_retry_propagates(self):
        initialize = Mock(side_effect=wandb.errors.CommError('unavailable'))
        with self.assertRaises(wandb.errors.CommError):
            with_wandb_offline_fallback(initialize)(frozen_config(), None)
        self.assertEqual(initialize.call_count, 2)


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