from __future__ import annotations

import gc
import os
import re
import time
from pathlib import Path

import inference.core.cache.air_gapped as air_gapped_module
import pytest
from inference.core.cache.air_gapped import (_slugify_model_id,
                                             is_block_cached, 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_list_format_returns_true_when_any_model_has_traditional_cache(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped_module, 'MODEL_CACHE_DIR', str(tmp_path))
    model_a = 'model_a'
    model_b = 'model_b'
    traditional_dir = tmp_path / model_b
    traditional_dir.mkdir()
    file_inside = traditional_dir / 'weights.bin'
    file_inside.write_text('dummy')
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_list_format_returns_true_when_any_model_has_traditional_cache', 'is_block_cached', '7', codeflash_loop_index, [model_a, model_b])
    codeflash_output = codeflash_return_value

def test_list_format_returns_false_when_no_models_cached(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped_module, 'MODEL_CACHE_DIR', str(tmp_path))
    model_x = 'missing_x'
    model_y = 'missing_y'
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_list_format_returns_false_when_no_models_cached', 'is_block_cached', '3', codeflash_loop_index, [model_x, model_y])
    codeflash_output = codeflash_return_value

def test_inference_models_layout_is_recognized_for_list_format(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped_module, 'MODEL_CACHE_DIR', str(tmp_path))
    model_id = 'Some.Model/With*Chars v1'
    slug = _slugify_model_id(model_id)
    models_cache_dir = tmp_path / 'models-cache' / slug
    models_cache_dir.mkdir(parents=True)
    (models_cache_dir / 'state.pt').write_text('data')
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_inference_models_layout_is_recognized_for_list_format', 'is_block_cached', '6', codeflash_loop_index, [model_id])
    codeflash_output = codeflash_return_value

def test_empty_list_is_treated_as_uncached(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped_module, 'MODEL_CACHE_DIR', str(tmp_path))
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_empty_list_is_treated_as_uncached', 'is_block_cached', '1', codeflash_loop_index, [])
    codeflash_output = codeflash_return_value

def test_non_list_non_dict_input_returns_false(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped_module, 'MODEL_CACHE_DIR', str(tmp_path))
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_non_list_non_dict_input_returns_false', 'is_block_cached', '1', codeflash_loop_index, None)
    codeflash_output = codeflash_return_value
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_non_list_non_dict_input_returns_false', 'is_block_cached', '2', codeflash_loop_index, 123)
    codeflash_output = codeflash_return_value
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_non_list_non_dict_input_returns_false', 'is_block_cached', '3', codeflash_loop_index, ('tuple', 'of', 'ids'))
    codeflash_output = codeflash_return_value

def test_legacy_dict_format_delegates_to_are_all_files_cached(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    artifacts_spec = {'model_id': 'legacy-model', 'files': ['a', 'b']}

    def fake_are_all_files_cached(files, model_id):
        return True
    monkeypatch.setattr(air_gapped_module, 'are_all_files_cached', fake_are_all_files_cached)
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_legacy_dict_format_delegates_to_are_all_files_cached', 'is_block_cached', '3', codeflash_loop_index, artifacts_spec)
    codeflash_output = codeflash_return_value

def test_legacy_dict_format_with_missing_keys_returns_false(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped_module, 'are_all_files_cached', lambda *a, **k: True)
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_legacy_dict_format_with_missing_keys_returns_false', 'is_block_cached', '1', codeflash_loop_index, {'files': ['a']})
    codeflash_output = codeflash_return_value
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_legacy_dict_format_with_missing_keys_returns_false', 'is_block_cached', '2', codeflash_loop_index, {'model_id': 'm', 'files': []})
    codeflash_output = codeflash_return_value
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_legacy_dict_format_with_missing_keys_returns_false', 'is_block_cached', '3', codeflash_loop_index, {'model_id': '', 'files': ['a']})
    codeflash_output = codeflash_return_value

def test_slug_truncation_and_special_char_handling(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped_module, 'MODEL_CACHE_DIR', str(tmp_path))
    long_model_id = 'x' * 200
    slug_long = _slugify_model_id(long_model_id)
    (tmp_path / 'models-cache' / slug_long).mkdir(parents=True)
    (tmp_path / 'models-cache' / slug_long / 'f').write_text('d')
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_slug_truncation_and_special_char_handling', 'is_block_cached', '5', codeflash_loop_index, [long_model_id])
    codeflash_output = codeflash_return_value
    special_only = '@@@!!!'
    slug_special = _slugify_model_id(special_only)
    (tmp_path / 'models-cache' / slug_special).mkdir(parents=True)
    (tmp_path / 'models-cache' / slug_special / 'f2').write_text('d')
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_slug_truncation_and_special_char_handling', 'is_block_cached', '10', codeflash_loop_index, [special_only])
    codeflash_output = codeflash_return_value

def test_large_list_of_model_ids_performance_and_correctness(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped_module, 'MODEL_CACHE_DIR', str(tmp_path))
    n = 1000
    ids = [f'm{i}' for i in range(n)]
    last_dir = tmp_path / ids[-1]
    last_dir.mkdir()
    (last_dir / 'file').write_text('x')
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_large_list_of_model_ids_performance_and_correctness', 'is_block_cached', '6', codeflash_loop_index, ids)
    codeflash_output = codeflash_return_value

def test_legacy_format_with_many_required_files_invokes_are_all_files_cached_once(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    files = [f'file_{i}.bin' for i in range(1000)]
    artifacts_spec = {'model_id': 'big-model', 'files': files}
    called = {'count': 0}

    def fake_are_all_files_cached(files_arg, model_id):
        called['count'] += 1
        return False
    monkeypatch.setattr(air_gapped_module, 'are_all_files_cached', fake_are_all_files_cached)
    codeflash_return_value = codeflash_wrap(is_block_cached, 'tests.inference.unit_tests.test_is_block_cached__unit_test_0', None, 'test_legacy_format_with_many_required_files_invokes_are_all_files_cached_once', 'is_block_cached', '5', codeflash_loop_index, artifacts_spec)
    codeflash_output = codeflash_return_value
    result = codeflash_output