from __future__ import annotations

import gc
import inspect
import os
import sys
import time
import types
from pathlib import Path

import dill as pickle
import inference.core.managers.decorators.fixed_size_cache as fixed_size_cache_module
import pytest
from inference.core.env import MEMORY_FREE_THRESHOLD, USE_INFERENCE_MODELS
from inference.core.managers.base import ModelManager
from inference.core.managers.decorators.fixed_size_cache import \
    WithFixedSizeCache
from inference.core.models.base import Model
from inference.core.registries.base import ModelRegistry


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}######!')
    iteration = os.environ['CODEFLASH_TEST_ITERATION']
    with Path('{codeflash_run_tmp_dir_client_side}', f'test_return_values_{iteration}.bin').open('ab') as f:
        pickled_values = pickle.dumps((args, kwargs, exception)) if exception else pickle.dumps((args, kwargs, return_value))
        _test_name = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_line_id}".encode('ascii')
        f.write(len(_test_name).to_bytes(4, byteorder='big'))
        f.write(_test_name)
        f.write(codeflash_duration.to_bytes(8, byteorder='big'))
        f.write(len(pickled_values).to_bytes(4, byteorder='big'))
        f.write(pickled_values)
        f.write(codeflash_loop_index.to_bytes(8, byteorder='big'))
        f.write(len(invocation_id).to_bytes(4, byteorder='big'))
        f.write(invocation_id.encode('ascii'))
    if exception:
        raise exception
    return return_value

def _install_fake_torch(is_available=True, mem_get_info_sequence=None):
    """
    Install a fake 'torch' module into sys.modules to control cuda behavior.

    Args:
        is_available (bool): what torch.cuda.is_available() should return.
        mem_get_info_sequence (list of (free, total) tuples or Exceptions): values
            mem_get_info() will return on successive calls. If an item is an Exception
            instance, mem_get_info will raise it.
    Returns:
        tuple(original_torch, state): original torch module (or None) and a dict
            containing counters we can inspect (e.g., empty_cache_calls).
    """
    original_torch = sys.modules.get('torch')
    fake_torch = types.ModuleType('torch')
    fake_cuda = types.ModuleType('torch.cuda')
    seq = list(mem_get_info_sequence) if mem_get_info_sequence is not None else []
    state = {'mem_get_info_calls': 0, 'empty_cache_calls': 0}

    def is_available_func():
        return bool(is_available)

    def mem_get_info_func():
        state['mem_get_info_calls'] += 1
        if not seq:
            return (8000, 16000)
        if len(seq) > 1:
            value = seq.pop(0)
        else:
            value = seq[0]
        if isinstance(value, Exception):
            raise value
        return value

    def empty_cache_func():
        state['empty_cache_calls'] += 1
    fake_cuda.is_available = is_available_func
    fake_cuda.mem_get_info = mem_get_info_func
    fake_cuda.empty_cache = empty_cache_func
    fake_torch.cuda = fake_cuda
    sys.modules['torch'] = fake_torch
    sys.modules['torch.cuda'] = fake_cuda
    return (original_torch, state)

def _restore_torch_module(original_torch):
    """
    Restore the original 'torch' module in sys.modules after a test.
    """
    sys.modules.pop('torch', None)
    sys.modules.pop('torch.cuda', None)
    if original_torch is not None:
        sys.modules['torch'] = original_torch

def test_returns_false_when_torch_missing():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original = sys.modules.pop('torch', None)
    try:
        mgr = ModelManager(ModelRegistry({}), models={})
        wrapper = WithFixedSizeCache(mgr, max_size=4)
        _call__bound__arguments = inspect.signature(wrapper.memory_pressure_detected).bind()
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(wrapper.memory_pressure_detected, 'tests.inference.unit_tests.test_memory_pressure_detected__unit_test_0', None, 'test_returns_false_when_torch_missing', 'memory_pressure_detected', '1_2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        if original is not None:
            sys.modules['torch'] = original

def test_returns_false_when_cuda_unavailable():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_torch, state = _install_fake_torch(is_available=False, mem_get_info_sequence=None)
    try:
        mgr = ModelManager(ModelRegistry({}), models={})
        wrapper = WithFixedSizeCache(mgr)
        _call__bound__arguments = inspect.signature(wrapper.memory_pressure_detected).bind()
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(wrapper.memory_pressure_detected, 'tests.inference.unit_tests.test_memory_pressure_detected__unit_test_0', None, 'test_returns_false_when_cuda_unavailable', 'memory_pressure_detected', '1_2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        _restore_torch_module(original_torch)

def test_memory_above_threshold_returns_false_no_empty_cache_called():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    threshold = fixed_size_cache_module.MEMORY_FREE_THRESHOLD
    total = 1000
    free = int((threshold + 0.05) * total)
    original_torch, state = _install_fake_torch(is_available=True, mem_get_info_sequence=[(free, total)])
    try:
        mgr = ModelManager(ModelRegistry({}), models={})
        wrapper = WithFixedSizeCache(mgr)
        prev_flag = getattr(fixed_size_cache_module, 'USE_INFERENCE_MODELS', None)
        fixed_size_cache_module.USE_INFERENCE_MODELS = True
        try:
            _call__bound__arguments = inspect.signature(wrapper.memory_pressure_detected).bind()
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(wrapper.memory_pressure_detected, 'tests.inference.unit_tests.test_memory_pressure_detected__unit_test_0', None, 'test_memory_above_threshold_returns_false_no_empty_cache_called', 'memory_pressure_detected', '4_4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
            codeflash_output = codeflash_return_value
            result = codeflash_output
        finally:
            fixed_size_cache_module.USE_INFERENCE_MODELS = prev_flag
    finally:
        _restore_torch_module(original_torch)

def test_memory_below_threshold_returns_true_when_use_inference_models_false():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    threshold = fixed_size_cache_module.MEMORY_FREE_THRESHOLD
    total = 1000
    free = int((threshold - 0.05) * total) if threshold > 0.06 else int(threshold / 2 * total)
    original_torch, state = _install_fake_torch(is_available=True, mem_get_info_sequence=[(free, total)])
    try:
        mgr = ModelManager(ModelRegistry({}), models={})
        wrapper = WithFixedSizeCache(mgr)
        prev_flag = getattr(fixed_size_cache_module, 'USE_INFERENCE_MODELS', None)
        fixed_size_cache_module.USE_INFERENCE_MODELS = False
        try:
            _call__bound__arguments = inspect.signature(wrapper.memory_pressure_detected).bind()
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(wrapper.memory_pressure_detected, 'tests.inference.unit_tests.test_memory_pressure_detected__unit_test_0', None, 'test_memory_below_threshold_returns_true_when_use_inference_models_false', 'memory_pressure_detected', '4_4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
            codeflash_output = codeflash_return_value
            result = codeflash_output
        finally:
            fixed_size_cache_module.USE_INFERENCE_MODELS = prev_flag
    finally:
        _restore_torch_module(original_torch)

def test_empty_cache_invoked_and_result_updated_after_empty_cache():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    threshold = fixed_size_cache_module.MEMORY_FREE_THRESHOLD
    total = 1000
    free_before = int((threshold - 0.05) * total) if threshold > 0.06 else int(threshold / 2 * total)
    free_after = int((threshold + 0.05) * total)
    original_torch, state = _install_fake_torch(is_available=True, mem_get_info_sequence=[(free_before, total), (free_after, total)])
    try:
        mgr = ModelManager(ModelRegistry({}), models={})
        wrapper = WithFixedSizeCache(mgr)
        prev_flag = getattr(fixed_size_cache_module, 'USE_INFERENCE_MODELS', None)
        fixed_size_cache_module.USE_INFERENCE_MODELS = True
        try:
            _call__bound__arguments = inspect.signature(wrapper.memory_pressure_detected).bind()
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(wrapper.memory_pressure_detected, 'tests.inference.unit_tests.test_memory_pressure_detected__unit_test_0', None, 'test_empty_cache_invoked_and_result_updated_after_empty_cache', 'memory_pressure_detected', '5_4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
            codeflash_output = codeflash_return_value
            result = codeflash_output
        finally:
            fixed_size_cache_module.USE_INFERENCE_MODELS = prev_flag
    finally:
        _restore_torch_module(original_torch)

def test_mem_get_info_exception_is_handled_and_returns_false():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_torch, state = _install_fake_torch(is_available=True, mem_get_info_sequence=[RuntimeError('cuda fail')])
    try:
        mgr = ModelManager(ModelRegistry({}), models={})
        wrapper = WithFixedSizeCache(mgr)
        _call__bound__arguments = inspect.signature(wrapper.memory_pressure_detected).bind()
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(wrapper.memory_pressure_detected, 'tests.inference.unit_tests.test_memory_pressure_detected__unit_test_0', None, 'test_mem_get_info_exception_is_handled_and_returns_false', 'memory_pressure_detected', '1_2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        _restore_torch_module(original_torch)

def test_division_by_zero_in_mem_get_info_is_handled_and_returns_false():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_torch, state = _install_fake_torch(is_available=True, mem_get_info_sequence=[(0, 0)])
    try:
        mgr = ModelManager(ModelRegistry({}), models={})
        wrapper = WithFixedSizeCache(mgr)
        _call__bound__arguments = inspect.signature(wrapper.memory_pressure_detected).bind()
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(wrapper.memory_pressure_detected, 'tests.inference.unit_tests.test_memory_pressure_detected__unit_test_0', None, 'test_division_by_zero_in_mem_get_info_is_handled_and_returns_false', 'memory_pressure_detected', '1_2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        _restore_torch_module(original_torch)

def test_large_scale_cache_initialization_and_memory_check():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    large_models = {f'm{i}': Model() for i in range(1000)}
    mgr = ModelManager(ModelRegistry({}), models=large_models)
    wrapper = WithFixedSizeCache(mgr, max_size=500)
    original = sys.modules.pop('torch', None)
    try:
        _call__bound__arguments = inspect.signature(wrapper.memory_pressure_detected).bind()
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(wrapper.memory_pressure_detected, 'tests.inference.unit_tests.test_memory_pressure_detected__unit_test_0', None, 'test_large_scale_cache_initialization_and_memory_check', 'memory_pressure_detected', '4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        if original is not None:
            sys.modules['torch'] = original