"""CUDA diagnostics and compiler propagation, without requiring a toolkit/GPU."""
import os
from pathlib import Path
import subprocess
import sys
import tempfile
from types import SimpleNamespace
import unittest
from unittest.mock import patch

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'tools'))
from check_cuda_toolchain import check_toolchain
from build_mfem_bundle import mpich_cuda_architectures, prepare_hypre_source


class CudaToolchainTest(unittest.TestCase):
    def test_probe_checks_both_compilers_without_running_gpu_code(self):
        calls = []

        def run(command, **kwargs):
            calls.append(command)
            if '--version' not in command:
                self.assertNotIn('NVCC_APPEND_FLAGS', kwargs['env'])
                self.assertNotIn('NVCC_CCBIN', kwargs['env'])
            return subprocess.CompletedProcess(command, 0, 'CUDA test toolkit', '')

        with patch.dict(os.environ, {'NVCC_APPEND_FLAGS': '-allow-unsupported-compiler',
                                     'NVCC_CCBIN': 'wrong-host'}), patch('subprocess.run', side_effect=run):
            check_toolchain('nvcc', 'gcc-compatible', 'g++-compatible')
        self.assertEqual(len(calls), 3)
        self.assertEqual(calls[1][2], 'gcc-compatible')
        self.assertIn('-c', calls[1])  # A C driver need not link the C++ runtime.
        self.assertEqual(calls[2][2], 'g++-compatible')
        self.assertNotIn('-c', calls[2])

    def test_mpich_architectures_follow_toolkit_and_user_selection(self):
        self.assertEqual(mpich_cuda_architectures('86-real;90-virtual;86', 'nvcc'), '86,90')
        self.assertEqual(mpich_cuda_architectures('native', 'nvcc'), 'auto')
        with patch('build_mfem_bundle.output', return_value='sm_75\nsm_80\nsm_86\nsm_90'):
            self.assertEqual(mpich_cuda_architectures('all-major', 'nvcc'), '75,80,90')
            self.assertEqual(mpich_cuda_architectures('all', 'nvcc'), '75,80,86,90')
        with self.assertRaisesRegex(RuntimeError, 'Invalid CUDA architecture'):
            mpich_cuda_architectures('garbage', 'nvcc')

    def test_hypre_patch_preserves_archive_and_older_cuda_branch(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            original = root / 'original'
            utilities = original / 'src' / 'utilities'
            utilities.mkdir(parents=True)
            header = ('#include <thrust/execution_policy.h>\n'
                      '#define HYPRE_THRUST_NOT(pred) thrust::not1(pred)\n'
                      'using namespace thrust::placeholders;\n')
            for name in ('device_utils.h', '_hypre_utilities.hpp'):
                (utilities / name).write_text(header)
            memory = ''.join(
                f'HYPRE_CUDA_CALL( cudaMemPrefetchAsync(ptr, size, {device},\n'
                '                                            hypre_HandleComputeStream(hypre_handle())) );\n'
                for device in ('hypre_HandleDevice(hypre_handle())', 'cudaCpuDeviceId'))
            (utilities / 'memory.c').write_text(memory)
            args = SimpleNamespace(hypre_source=original, work_dir=root / 'work')
            self.assertEqual(prepare_hypre_source(args, {'cuda': False}), original)
            patched = prepare_hypre_source(args, {'cuda': True})
            self.assertEqual((utilities / 'memory.c').read_text(), memory)
            updated = (patched / 'src/utilities/memory.c').read_text()
            self.assertIn('#if CUDART_VERSION >= 13000', updated)
            self.assertIn('cudaMemLocationTypeHost, 0', updated)
            self.assertIn('cudaMemPrefetchAsync(ptr, size, cudaCpuDeviceId,', updated)
            self.assertIn('#define HYPRE_THRUST_NOT(pred) thrust::not1(pred)',
                          (patched / 'src/utilities/device_utils.h').read_text())
            prepare_hypre_source(args, {'cuda': True})
            self.assertEqual((patched / 'src/utilities/memory.c').read_text(), updated)

    def test_cxx_failure_reports_toolkit_host_and_recovery(self):
        results = [subprocess.CompletedProcess([], 0, 'CUDA test toolkit', ''),
                   subprocess.CompletedProcess([], 0, '', ''),
                   subprocess.CompletedProcess([], 1, '', 'unsupported Microsoft Visual Studio version')]
        with patch('subprocess.run', side_effect=results):
            with self.assertRaises(RuntimeError) as caught:
                check_toolchain('nvcc.exe', 'cl.exe', 'cl.exe')
        message = str(caught.exception)
        for expected in ['C++ compiler: cl.exe', 'CUDA test toolkit',
                         'unsupported Microsoft Visual Studio version', 'fresh Meson build directory']:
            self.assertIn(expected, message)


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