from __future__ import annotations

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

import dill as pickle
import pytest
from inference.core.cache import air_gapped
from inference.core.cache.air_gapped import is_model_cached
from inference.core.env import MODEL_CACHE_DIR, USE_INFERENCE_MODELS


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
is_model_cached = air_gapped.is_model_cached
_slugify_model_id = air_gapped._slugify_model_id
_has_non_hidden_children = air_gapped._has_non_hidden_children

def _write_file(path: Path, name: str='file.bin', hidden: bool=False):
    """Helper to create a file inside path. Use hidden=True to create a dot-file."""
    fname = f'.{name}' if hidden else name
    p = path / fname
    p.write_text('data')
    return p

def test_traditional_cache_detected_when_inference_models_disabled(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    air_gapped.USE_INFERENCE_MODELS = False
    model_id = 'my-model'
    traditional_dir = tmp_path / model_id
    traditional_dir.mkdir()
    _write_file(traditional_dir, 'weights.bin', hidden=False)
    _call__bound__arguments = inspect.signature(is_model_cached).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_0', None, 'test_traditional_cache_detected_when_inference_models_disabled', 'is_model_cached', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_traditional_hidden_files_are_ignored(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    air_gapped.USE_INFERENCE_MODELS = False
    model_id = 'hidden-only-model'
    d = tmp_path / model_id
    d.mkdir()
    _write_file(d, 'lock', hidden=True)
    _call__bound__arguments = inspect.signature(is_model_cached).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_0', None, 'test_traditional_hidden_files_are_ignored', 'is_model_cached', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_missing_traditional_directory_returns_false(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    air_gapped.USE_INFERENCE_MODELS = False
    _call__bound__arguments = inspect.signature(is_model_cached).bind('nonexistent-model')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_0', None, 'test_missing_traditional_directory_returns_false', 'is_model_cached', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_inference_models_enabled_prefers_traditional_before_models_cache(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    air_gapped.USE_INFERENCE_MODELS = True
    model_id = 'migrated-model'
    traditional_dir = tmp_path / model_id
    traditional_dir.mkdir()
    _write_file(traditional_dir, 'artifact.bin', hidden=False)
    slug = _slugify_model_id(model_id)
    (tmp_path / 'models-cache' / slug).mkdir(parents=True, exist_ok=True)
    _call__bound__arguments = inspect.signature(is_model_cached).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_0', None, 'test_inference_models_enabled_prefers_traditional_before_models_cache', 'is_model_cached', '8', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_inference_models_enabled_detects_models_cache_layout(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    air_gapped.USE_INFERENCE_MODELS = True
    model_id = 'new-layout-model!@#$%^&*()'
    slug = _slugify_model_id(model_id)
    models_cache_dir = tmp_path / 'models-cache' / slug
    models_cache_dir.mkdir(parents=True)
    _write_file(models_cache_dir, 'model.pb', hidden=False)
    _call__bound__arguments = inspect.signature(is_model_cached).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_0', None, 'test_inference_models_enabled_detects_models_cache_layout', 'is_model_cached', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_slugify_truncation_and_digest_produces_expected_directory_name():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    long_id = '!' * 60 + 'example-model-with-ALPHA_and_123'
    slug = _slugify_model_id(long_id)
    base, digest = slug.rsplit('-', 1)

def test_has_non_hidden_children_handles_non_directory(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    f = tmp_path / 'a_regular_file.txt'
    f.write_text('content')

def test_none_model_id_raises_type_error(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    air_gapped.USE_INFERENCE_MODELS = True
    with pytest.raises(TypeError):
        _call__bound__arguments = inspect.signature(is_model_cached).bind(None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(is_model_cached, 'tests.inference.unit_tests.test_is_model_cached__unit_test_0', None, 'test_none_model_id_raises_type_error', 'is_model_cached', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_number_of_traditional_directories(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    air_gapped.MODEL_CACHE_DIR = str(tmp_path)
    air_gapped.USE_INFERENCE_MODELS = False
    count = 1000
    target_index = 777
    for i in range(count):
        d = tmp_path / f'model_{i}'
        d.mkdir()
        if i == target_index:
            _write_file(d, 'weights.bin', hidden=False)
        else:
            _write_file(d, 'lock', hidden=True)
    positives = [is_model_cached(f'model_{i}') for i in range(count)]