from __future__ import annotations

import gc
import os
import shutil
import tempfile
import time
from unittest import mock

import pytest
from inference.core.cache.air_gapped import (_has_non_hidden_children,
                                             _slugify_model_id,
                                             is_model_cached)


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 test_model_not_cached_when_traditional_path_does_not_exist():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "Test that is_model_cached returns False when model_id path doesn't exist."
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', False):
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_not_cached_when_traditional_path_does_not_exist', 'is_model_cached', '1_0_0_0', codeflash_loop_index, 'nonexistent-model-id')
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_model_cached_in_traditional_layout_with_use_inference_models_false():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that model is found in traditional layout when USE_INFERENCE_MODELS is False.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', False):
                model_id = 'bert-base-uncased'
                model_path = os.path.join(temp_dir, model_id)
                os.makedirs(model_path)
                with open(os.path.join(model_path, 'config.json'), 'w') as f:
                    f.write('{}')
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_cached_in_traditional_layout_with_use_inference_models_false', 'is_model_cached', '1_0_0_4', codeflash_loop_index, model_id)
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_model_not_cached_when_traditional_directory_is_empty():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that is_model_cached returns False when traditional directory is empty.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', False):
                model_id = 'empty-model'
                model_path = os.path.join(temp_dir, model_id)
                os.makedirs(model_path)
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_not_cached_when_traditional_directory_is_empty', 'is_model_cached', '1_0_0_3', codeflash_loop_index, model_id)
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_model_cached_in_inference_models_layout():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that model is found in inference-models layout when USE_INFERENCE_MODELS is True.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', True):
                model_id = 'yolov8n'
                slug = _slugify_model_id(model_id)
                models_cache_path = os.path.join(temp_dir, 'models-cache', slug)
                os.makedirs(models_cache_path)
                with open(os.path.join(models_cache_path, 'model.pt'), 'w') as f:
                    f.write('model data')
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_cached_in_inference_models_layout', 'is_model_cached', '1_0_0_5', codeflash_loop_index, model_id)
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_model_cached_in_both_layouts_returns_true():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that model is found when it exists in both traditional and inference-models layouts.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', True):
                model_id = 'bert-large-uncased'
                traditional_path = os.path.join(temp_dir, model_id)
                os.makedirs(traditional_path)
                with open(os.path.join(traditional_path, 'config.json'), 'w') as f:
                    f.write('{}')
                slug = _slugify_model_id(model_id)
                models_cache_path = os.path.join(temp_dir, 'models-cache', slug)
                os.makedirs(models_cache_path)
                with open(os.path.join(models_cache_path, 'model.pt'), 'w') as f:
                    f.write('model data')
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_cached_in_both_layouts_returns_true', 'is_model_cached', '1_0_0_8', codeflash_loop_index, model_id)
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_model_found_in_traditional_when_inference_models_enabled():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that traditional layout is checked first when USE_INFERENCE_MODELS is True.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', True):
                model_id = 'roberta-base'
                traditional_path = os.path.join(temp_dir, model_id)
                os.makedirs(traditional_path)
                with open(os.path.join(traditional_path, 'pytorch_model.bin'), 'w') as f:
                    f.write('model data')
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_found_in_traditional_when_inference_models_enabled', 'is_model_cached', '1_0_0_4', codeflash_loop_index, model_id)
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_has_non_hidden_children_with_single_hidden_file():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that _has_non_hidden_children returns False with only hidden files.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with open(os.path.join(temp_dir, '.lock'), 'w') as f:
            f.write('')
        result = _has_non_hidden_children(temp_dir)

def test_has_non_hidden_children_with_mixed_files():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that _has_non_hidden_children returns True with both hidden and non-hidden files.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with open(os.path.join(temp_dir, '.lock'), 'w') as f:
            f.write('')
        with open(os.path.join(temp_dir, 'data.txt'), 'w') as f:
            f.write('data')
        result = _has_non_hidden_children(temp_dir)

def test_model_id_with_special_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that model_id with special characters is handled correctly.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', False):
                model_id = 'model/with-special@chars!'
                traditional_path = os.path.join(temp_dir, model_id)
                os.makedirs(traditional_path, exist_ok=True)
                with open(os.path.join(traditional_path, 'config.json'), 'w') as f:
                    f.write('{}')
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_id_with_special_characters', 'is_model_cached', '1_0_0_4', codeflash_loop_index, model_id)
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_slugify_model_id_with_special_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test _slugify_model_id converts special characters to hyphens.'
    model_id = 'model@@@name///test'
    slug = _slugify_model_id(model_id)

def test_slugify_model_id_with_consecutive_separators():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test _slugify_model_id collapses consecutive underscores and hyphens.'
    model_id = 'model__name--test'
    slug = _slugify_model_id(model_id)

def test_slugify_model_id_with_only_special_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test _slugify_model_id handles model ID with only special characters.'
    model_id = '@@@###$$$'
    slug = _slugify_model_id(model_id)

def test_slugify_model_id_truncates_long_names():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test _slugify_model_id truncates long model names.'
    long_model_id = 'a' * 100
    slug = _slugify_model_id(long_model_id)
    slug_parts = slug.rsplit('-', 1)
    slug_name = slug_parts[0]

def test_model_id_empty_string():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test is_model_cached with empty string model_id.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', False):
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_id_empty_string', 'is_model_cached', '1_0_0_0', codeflash_loop_index, '')
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_has_non_hidden_children_with_nonexistent_path():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that _has_non_hidden_children returns False for nonexistent paths.'
    result = _has_non_hidden_children('/nonexistent/path/that/does/not/exist')

def test_has_non_hidden_children_with_permission_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that _has_non_hidden_children handles permission errors gracefully.'
    with tempfile.TemporaryDirectory() as temp_dir:
        restricted_dir = os.path.join(temp_dir, 'restricted')
        os.makedirs(restricted_dir)
        os.chmod(restricted_dir, 0)
        try:
            result = _has_non_hidden_children(restricted_dir)
        finally:
            os.chmod(restricted_dir, 493)

def test_model_directory_with_only_dot_files():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that directory with only dotfiles is considered empty.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', False):
                model_id = 'test-model'
                model_path = os.path.join(temp_dir, model_id)
                os.makedirs(model_path)
                with open(os.path.join(model_path, '.gitkeep'), 'w') as f:
                    f.write('')
                with open(os.path.join(model_path, '.lock'), 'w') as f:
                    f.write('')
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_directory_with_only_dot_files', 'is_model_cached', '1_0_0_5', codeflash_loop_index, model_id)
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_model_id_with_path_traversal_attempt():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that model_id with path traversal sequences is handled safely.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', False):
                model_id = '../../../etc/passwd'
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_id_with_path_traversal_attempt', 'is_model_cached', '1_0_0_1', codeflash_loop_index, model_id)
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_slugify_produces_consistent_hash():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that _slugify_model_id produces consistent hash for same input.'
    model_id = 'consistent-test-model'
    slug1 = _slugify_model_id(model_id)
    slug2 = _slugify_model_id(model_id)

def test_slugify_produces_different_hash_for_different_input():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that _slugify_model_id produces different hash for different inputs.'
    slug1 = _slugify_model_id('model-one')
    slug2 = _slugify_model_id('model-two')

def test_model_cached_with_many_files_in_directory():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that is_model_cached works efficiently with many files in model directory.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', False):
                model_id = 'large-model'
                model_path = os.path.join(temp_dir, model_id)
                os.makedirs(model_path)
                for i in range(100):
                    with open(os.path.join(model_path, f'file_{i}.dat'), 'w') as f:
                        f.write(f'data {i}')
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_cached_with_many_files_in_directory', 'is_model_cached', '1_0_0_4', codeflash_loop_index, model_id)
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_model_cached_with_many_hidden_files():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that is_model_cached correctly handles directories with many hidden files.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', False):
                model_id = 'hidden-files-model'
                model_path = os.path.join(temp_dir, model_id)
                os.makedirs(model_path)
                for i in range(100):
                    with open(os.path.join(model_path, f'.hidden_{i}'), 'w') as f:
                        f.write('')
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_cached_with_many_hidden_files', 'is_model_cached', '1_0_0_4', codeflash_loop_index, model_id)
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_model_cached_with_mixed_files_large_scale():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that is_model_cached efficiently finds non-hidden files among many files.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', False):
                model_id = 'mixed-files-model'
                model_path = os.path.join(temp_dir, model_id)
                os.makedirs(model_path)
                for i in range(500):
                    with open(os.path.join(model_path, f'.hidden_{i}'), 'w') as f:
                        f.write('')
                for i in range(10):
                    with open(os.path.join(model_path, f'visible_{i}.dat'), 'w') as f:
                        f.write(f'data {i}')
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_cached_with_mixed_files_large_scale', 'is_model_cached', '1_0_0_5', codeflash_loop_index, model_id)
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_model_cached_checks_both_layouts_large_scale():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that is_model_cached efficiently checks both layouts with many models.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', True):
                for i in range(50):
                    model_id = f'model-{i}'
                    slug = _slugify_model_id(model_id)
                    models_cache_path = os.path.join(temp_dir, 'models-cache', slug)
                    os.makedirs(models_cache_path)
                    with open(os.path.join(models_cache_path, 'model.pt'), 'w') as f:
                        f.write('data')
                codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_model_cached_checks_both_layouts_large_scale', 'is_model_cached', '1_0_0_1', codeflash_loop_index, 'model-25')
                codeflash_output = codeflash_return_value
                result = codeflash_output

def test_has_non_hidden_children_with_many_entries():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that _has_non_hidden_children efficiently checks large directories.'
    with tempfile.TemporaryDirectory() as temp_dir:
        for i in range(1000):
            with open(os.path.join(temp_dir, f'file_{i:04d}.dat'), 'w') as f:
                f.write(f'data {i}')
        result = _has_non_hidden_children(temp_dir)

def test_multiple_model_cache_checks_in_sequence():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that multiple cache checks work correctly in sequence.'
    with tempfile.TemporaryDirectory() as temp_dir:
        with mock.patch('inference.core.cache.air_gapped.MODEL_CACHE_DIR', temp_dir):
            with mock.patch('inference.core.cache.air_gapped.USE_INFERENCE_MODELS', False):
                for i in range(20):
                    model_id = f'seq-model-{i}'
                    model_path = os.path.join(temp_dir, model_id)
                    os.makedirs(model_path)
                    if i % 2 == 0:
                        with open(os.path.join(model_path, 'config.json'), 'w') as f:
                            f.write('{}')
                results = []
                for i in range(20):
                    model_id = f'seq-model-{i}'
                    codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_1', None, 'test_multiple_model_cache_checks_in_sequence', 'is_model_cached', '1_0_0_2_1', codeflash_loop_index, model_id)
                for i, result in enumerate(results):
                    pass

def test_slugify_with_batch_of_model_ids():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that _slugify_model_id works correctly on a batch of diverse model IDs.'
    model_ids = ['bert-base-uncased', 'gpt2', 'resnet50', 'vit-base-patch16-224', 'roberta-large', 'albert-xxlarge-v1', 'electra-small-discriminator', 'xlnet-base-cased', 't5-small', 'bart-base']
    slugs = [_slugify_model_id(model_id) for model_id in model_ids]