from __future__ import annotations

import builtins
import gc
import os
import sys
import time
import types

import inference.core.env as env_module
import inference.core.logger as logger_module
import pytest
from inference.core.env import USE_INFERENCE_MODELS
from inference.core.managers.base import try_releasing_cuda_memory


def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, *args, **kwargs):
    test_id = f'{codeflash_test_module_name}:{codeflash_test_class_name}:{codeflash_test_name}:{codeflash_line_id}:{codeflash_loop_index}'
    if not hasattr(codeflash_wrap, 'index'):
        codeflash_wrap.index = {}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f'{codeflash_line_id}_{codeflash_test_index}'
    test_stdout_tag = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_loop_index}:{invocation_id}"
    print(f'!$######{test_stdout_tag}######$!')
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        exception = e
    gc.enable()
    print(f'!######{test_stdout_tag}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

def _make_fake_torch(is_available=True, empty_cache_raises=False, empty_cache_side_effect=None):
    """
    Construct a real ModuleType to act as a fake 'torch' package.
    - is_available: value returned by torch.cuda.is_available()
    - empty_cache_raises: if True, torch.cuda.empty_cache will raise an Exception
    - empty_cache_side_effect: optional callable to be executed when empty_cache is called
    Returns the module object and a dict with counters we can inspect.
    """
    counters = {'empty_cache_calls': 0, 'is_available_calls': 0}
    torch_mod = types.ModuleType('torch')
    cuda_mod = types.ModuleType('torch.cuda')

    def is_available_fn():
        counters['is_available_calls'] += 1
        return is_available

    def empty_cache_fn():
        counters['empty_cache_calls'] += 1
        if empty_cache_side_effect is not None:
            empty_cache_side_effect()
        if empty_cache_raises:
            raise RuntimeError('simulated empty_cache failure')
    cuda_mod.is_available = is_available_fn
    cuda_mod.empty_cache = empty_cache_fn
    torch_mod.cuda = cuda_mod
    return (torch_mod, counters)

def test_returns_none_when_use_inference_models_false():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_flag = env_module.USE_INFERENCE_MODELS
    try:
        env_module.USE_INFERENCE_MODELS = False
        codeflash_return_value = codeflash_wrap(try_releasing_cuda_memory, 'tests.inference.unit_tests.test_try_releasing_cuda_memory__unit_test_0', None, 'test_returns_none_when_use_inference_models_false', 'try_releasing_cuda_memory', '1_1', codeflash_loop_index)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        env_module.USE_INFERENCE_MODELS = original_flag

def test_returns_none_when_use_inference_models_none():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_flag = env_module.USE_INFERENCE_MODELS
    try:
        env_module.USE_INFERENCE_MODELS = None
        codeflash_return_value = codeflash_wrap(try_releasing_cuda_memory, 'tests.inference.unit_tests.test_try_releasing_cuda_memory__unit_test_0', None, 'test_returns_none_when_use_inference_models_none', 'try_releasing_cuda_memory', '1_1', codeflash_loop_index)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        env_module.USE_INFERENCE_MODELS = original_flag

def test_calls_empty_cache_when_cuda_available_and_torch_present():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_flag = env_module.USE_INFERENCE_MODELS
    original_torch = sys.modules.get('torch')
    try:
        env_module.USE_INFERENCE_MODELS = True
        fake_torch, counters = _make_fake_torch(is_available=True)
        sys.modules['torch'] = fake_torch
        codeflash_return_value = codeflash_wrap(try_releasing_cuda_memory, 'tests.inference.unit_tests.test_try_releasing_cuda_memory__unit_test_0', None, 'test_calls_empty_cache_when_cuda_available_and_torch_present', 'try_releasing_cuda_memory', '2_3', codeflash_loop_index)
    finally:
        if original_torch is not None:
            sys.modules['torch'] = original_torch
        else:
            sys.modules.pop('torch', None)
        env_module.USE_INFERENCE_MODELS = original_flag

def test_does_not_call_empty_cache_when_cuda_not_available():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_flag = env_module.USE_INFERENCE_MODELS
    original_torch = sys.modules.get('torch')
    try:
        env_module.USE_INFERENCE_MODELS = True
        fake_torch, counters = _make_fake_torch(is_available=False)
        sys.modules['torch'] = fake_torch
        codeflash_return_value = codeflash_wrap(try_releasing_cuda_memory, 'tests.inference.unit_tests.test_try_releasing_cuda_memory__unit_test_0', None, 'test_does_not_call_empty_cache_when_cuda_not_available', 'try_releasing_cuda_memory', '2_3', codeflash_loop_index)
    finally:
        if original_torch is not None:
            sys.modules['torch'] = original_torch
        else:
            sys.modules.pop('torch', None)
        env_module.USE_INFERENCE_MODELS = original_flag

def test_importerror_when_importing_torch_is_handled_gracefully():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_flag = env_module.USE_INFERENCE_MODELS
    original_import = builtins.__import__
    original_torch = sys.modules.get('torch')
    try:
        env_module.USE_INFERENCE_MODELS = True
        sys.modules.pop('torch', None)

        def import_raising(name, globals=None, locals=None, fromlist=(), level=0):
            if name == 'torch' or (isinstance(name, str) and name.startswith('torch.')):
                raise ImportError('simulated missing torch')
            return original_import(name, globals, locals, fromlist, level)
        builtins.__import__ = import_raising
        codeflash_return_value = codeflash_wrap(try_releasing_cuda_memory, 'tests.inference.unit_tests.test_try_releasing_cuda_memory__unit_test_0', None, 'test_importerror_when_importing_torch_is_handled_gracefully', 'try_releasing_cuda_memory', '3_4', codeflash_loop_index)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        builtins.__import__ = original_import
        if original_torch is not None:
            sys.modules['torch'] = original_torch
        else:
            sys.modules.pop('torch', None)
        env_module.USE_INFERENCE_MODELS = original_flag

def test_logs_warning_when_empty_cache_raises(caplog):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_flag = env_module.USE_INFERENCE_MODELS
    original_torch = sys.modules.get('torch')
    try:
        env_module.USE_INFERENCE_MODELS = True
        fake_torch, counters = _make_fake_torch(is_available=True, empty_cache_raises=True)
        sys.modules['torch'] = fake_torch
        caplog.clear()
        codeflash_return_value = codeflash_wrap(try_releasing_cuda_memory, 'tests.inference.unit_tests.test_try_releasing_cuda_memory__unit_test_0', None, 'test_logs_warning_when_empty_cache_raises', 'try_releasing_cuda_memory', '2_4', codeflash_loop_index)
        warnings = [rec for rec in caplog.records if rec.levelname == 'WARNING']
    finally:
        if original_torch is not None:
            sys.modules['torch'] = original_torch
        else:
            sys.modules.pop('torch', None)
        env_module.USE_INFERENCE_MODELS = original_flag

def test_repeated_calls_empty_cache_many_times():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_flag = env_module.USE_INFERENCE_MODELS
    original_torch = sys.modules.get('torch')
    try:
        env_module.USE_INFERENCE_MODELS = True
        fake_torch, counters = _make_fake_torch(is_available=True)
        sys.modules['torch'] = fake_torch
        for _ in range(1000):
            codeflash_return_value = codeflash_wrap(try_releasing_cuda_memory, 'tests.inference.unit_tests.test_try_releasing_cuda_memory__unit_test_0', None, 'test_repeated_calls_empty_cache_many_times', 'try_releasing_cuda_memory', '2_3_0', codeflash_loop_index)
    finally:
        if original_torch is not None:
            sys.modules['torch'] = original_torch
        else:
            sys.modules.pop('torch', None)
        env_module.USE_INFERENCE_MODELS = original_flag