from __future__ import annotations

import gc
import hashlib
import importlib
import os
import pathlib
import time

import pytest
from inference.core.cache.air_gapped import _is_model_cached, _slugify_model_id
from inference.core.env import MODEL_CACHE_DIR


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
air_gapped = importlib.import_module('inference.core.cache.air_gapped')

def test_traditional_layout_detects_cache_basic(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    model_dir = tmp_path / 'modelA'
    model_dir.mkdir(parents=True, exist_ok=True)
    (model_dir / 'weights.bin').write_bytes(b'binaryweights')
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_traditional_layout_detects_cache_basic', '_is_model_cached', '4', codeflash_loop_index, 'modelA')
    codeflash_output = codeflash_return_value
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_traditional_layout_detects_cache_basic', '_is_model_cached', '5', codeflash_loop_index, 'no-such-model')
    codeflash_output = codeflash_return_value

def test_inference_models_layout_detects_cache_with_subdirectory(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    model_id = 'My Model@@@/v1'
    slug = _slugify_model_id(model_id)
    slug_dir = tmp_path / 'models-cache' / slug
    (slug_dir / 'subdir').mkdir(parents=True, exist_ok=True)
    (slug_dir / 'subdir' / 'config.json').write_text('{"ok": true}')
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_inference_models_layout_detects_cache_with_subdirectory', '_is_model_cached', '6', codeflash_loop_index, model_id)
    codeflash_output = codeflash_return_value

def test_empty_directories_do_not_count_as_cached(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    (tmp_path / 'empty_model').mkdir(parents=True, exist_ok=True)
    empty_slug = _slugify_model_id('++++')
    (tmp_path / 'models-cache' / empty_slug).mkdir(parents=True, exist_ok=True)
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_empty_directories_do_not_count_as_cached', '_is_model_cached', '4', codeflash_loop_index, 'empty_model')
    codeflash_output = codeflash_return_value
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_empty_directories_do_not_count_as_cached', '_is_model_cached', '5', codeflash_loop_index, '++++')
    codeflash_output = codeflash_return_value

def test_slugify_edge_cases_and_truncation_behavior(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    long_model_id = 'a' * 100
    long_slug = _slugify_model_id(long_model_id)
    (tmp_path / 'models-cache' / long_slug).mkdir(parents=True, exist_ok=True)
    (tmp_path / 'models-cache' / long_slug / 'file.bin').write_bytes(b'x')
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_slugify_edge_cases_and_truncation_behavior', '_is_model_cached', '5', codeflash_loop_index, long_model_id)
    codeflash_output = codeflash_return_value
    special_id = '@@@'
    special_slug = _slugify_model_id(special_id)
    (tmp_path / 'models-cache' / special_slug).mkdir(parents=True, exist_ok=True)
    (tmp_path / 'models-cache' / special_slug / 'f').write_text('1')
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_slugify_edge_cases_and_truncation_behavior', '_is_model_cached', '10', codeflash_loop_index, special_id)
    codeflash_output = codeflash_return_value

def test_invalid_type_raises_type_error(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    with pytest.raises(TypeError):
        codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_invalid_type_raises_type_error', '_is_model_cached', '1_0', codeflash_loop_index, None)

def test_large_scale_many_model_directories(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    n = 1000
    base = tmp_path
    for i in range(n):
        name = f'm{i:04d}'
        d = base / name
        d.mkdir(parents=True, exist_ok=True)
        (d / 'f.txt').write_text('x')
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_large_scale_many_model_directories', '_is_model_cached', '4', codeflash_loop_index, 'm0000')
    codeflash_output = codeflash_return_value
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_large_scale_many_model_directories', '_is_model_cached', '5', codeflash_loop_index, f'm{n - 1:04d}')
    codeflash_output = codeflash_return_value
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_large_scale_many_model_directories', '_is_model_cached', '6', codeflash_loop_index, 'm0500')
    codeflash_output = codeflash_return_value
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_large_scale_many_model_directories', '_is_model_cached', '7', codeflash_loop_index, 'm1001')
    codeflash_output = codeflash_return_value

def test_models_cache_single_directory_with_many_entries(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    model_id = 'heavy-model'
    slug = _slugify_model_id(model_id)
    slug_dir = tmp_path / 'models-cache' / slug
    slug_dir.mkdir(parents=True, exist_ok=True)
    for i in range(1000):
        (slug_dir / f'file_{i:04d}.dat').write_text('data')
    codeflash_return_value = codeflash_wrap(_is_model_cached, 'tests.inference.unit_tests.test__is_model_cached__unit_test_0', None, 'test_models_cache_single_directory_with_many_entries', '_is_model_cached', '6', codeflash_loop_index, model_id)
    codeflash_output = codeflash_return_value