"""Tests para torch_device_select."""

import sys
import unittest
import warnings

sys.path.insert(0, "python/functions")

import ml.torch_device_select as tds_module
from ml.torch_device_select import torch_device_select


class TestTorchDeviceSelect(unittest.TestCase):

    def _patch(self, cuda=False, mps=False, cuda_count=0):
        """Helper: parchea los helpers internos del modulo."""
        from unittest.mock import patch
        return [
            patch.object(tds_module, "_cuda_available", return_value=cuda),
            patch.object(tds_module, "_mps_available", return_value=mps),
            patch.object(tds_module, "_cuda_device_count", return_value=cuda_count),
        ]

    def test_preference_cpu_siempre_retorna_cpu(self):
        """preference=cpu siempre retorna cpu"""
        self.assertEqual(torch_device_select("cpu"), "cpu")

    def test_preference_auto_sin_cuda_ni_mps_retorna_cpu(self):
        """preference=auto sin cuda ni mps retorna cpu"""
        patches = self._patch(cuda=False, mps=False)
        for p in patches:
            p.start()
        try:
            self.assertEqual(torch_device_select("auto"), "cpu")
        finally:
            for p in patches:
                p.stop()

    def test_preference_cuda_sin_cuda_disponible_retorna_cpu_con_warning(self):
        """preference=cuda sin cuda disponible retorna cpu con warning"""
        patches = self._patch(cuda=False, mps=False)
        for p in patches:
            p.start()
        try:
            with warnings.catch_warnings(record=True) as w:
                warnings.simplefilter("always")
                result = torch_device_select("cuda")
            self.assertEqual(result, "cpu")
            self.assertTrue(any("CUDA" in str(warning.message) for warning in w))
        finally:
            for p in patches:
                p.stop()

    def test_preference_cuda_N_con_solo_1_GPU_retorna_cpu_con_warning(self):
        """preference=cuda:5 con solo 1 GPU retorna cpu con warning"""
        patches = self._patch(cuda=True, cuda_count=1)
        for p in patches:
            p.start()
        try:
            with warnings.catch_warnings(record=True) as w:
                warnings.simplefilter("always")
                result = torch_device_select("cuda:5")
            self.assertEqual(result, "cpu")
            self.assertTrue(len(w) > 0)
        finally:
            for p in patches:
                p.stop()

    def test_preference_desconocida_retorna_cpu_con_warning(self):
        """preference desconocida retorna cpu con warning"""
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            result = torch_device_select("vulkan")
        self.assertEqual(result, "cpu")
        self.assertTrue(len(w) > 0)


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