from __future__ import annotations

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

import dill as pickle
import pytest
from inference.core.cache.air_gapped import _slugify_model_id


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 test_basic_alphanumeric_preserved_and_digest_appended():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_id = 'Model123'
    expected_digest = hashlib.blake2s(model_id.encode('utf-8'), digest_size=4).hexdigest()
    _call__bound__arguments = inspect.signature(_slugify_model_id).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_slugify_model_id, 'tests.inference.unit_tests.test__slugify_model_id__unit_test_0', None, 'test_basic_alphanumeric_preserved_and_digest_appended', '_slugify_model_id', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_case_is_preserved_and_underscores_kept():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_id = 'AbC-Def_123'
    expected_digest = hashlib.blake2s(model_id.encode('utf-8'), digest_size=4).hexdigest()
    _call__bound__arguments = inspect.signature(_slugify_model_id).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_slugify_model_id, 'tests.inference.unit_tests.test__slugify_model_id__unit_test_0', None, 'test_case_is_preserved_and_underscores_kept', '_slugify_model_id', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_empty_input_uses_special_marker_and_digest_of_empty_string():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_id = ''
    expected_digest = hashlib.blake2s(model_id.encode('utf-8'), digest_size=4).hexdigest()
    _call__bound__arguments = inspect.signature(_slugify_model_id).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_slugify_model_id, 'tests.inference.unit_tests.test__slugify_model_id__unit_test_0', None, 'test_empty_input_uses_special_marker_and_digest_of_empty_string', '_slugify_model_id', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_only_special_characters_produces_single_hyphen_slug_then_digest():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_id = '@@@!!!'
    expected_digest = hashlib.blake2s(model_id.encode('utf-8'), digest_size=4).hexdigest()
    _call__bound__arguments = inspect.signature(_slugify_model_id).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_slugify_model_id, 'tests.inference.unit_tests.test__slugify_model_id__unit_test_0', None, 'test_only_special_characters_produces_single_hyphen_slug_then_digest', '_slugify_model_id', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_consecutive_underscores_and_dashes_are_consolidated():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_id = 'a__--__b'
    expected_digest = hashlib.blake2s(model_id.encode('utf-8'), digest_size=4).hexdigest()
    _call__bound__arguments = inspect.signature(_slugify_model_id).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_slugify_model_id, 'tests.inference.unit_tests.test__slugify_model_id__unit_test_0', None, 'test_consecutive_underscores_and_dashes_are_consolidated', '_slugify_model_id', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_special_character_sequences_replaced_and_consolidated():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_id = 'Model@@@Name__--X'
    expected_digest = hashlib.blake2s(model_id.encode('utf-8'), digest_size=4).hexdigest()
    _call__bound__arguments = inspect.signature(_slugify_model_id).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_slugify_model_id, 'tests.inference.unit_tests.test__slugify_model_id__unit_test_0', None, 'test_special_character_sequences_replaced_and_consolidated', '_slugify_model_id', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_unicode_characters_are_replaced_with_hyphen_preserving_allowed_tail():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_id = '模型😊ID'
    expected_digest = hashlib.blake2s(model_id.encode('utf-8'), digest_size=4).hexdigest()
    _call__bound__arguments = inspect.signature(_slugify_model_id).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_slugify_model_id, 'tests.inference.unit_tests.test__slugify_model_id__unit_test_0', None, 'test_unicode_characters_are_replaced_with_hyphen_preserving_allowed_tail', '_slugify_model_id', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_length_truncation_uses_original_for_digest():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_id = 'A' * 100
    expected_slug_part = 'A' * 48
    expected_digest = hashlib.blake2s(model_id.encode('utf-8'), digest_size=4).hexdigest()
    _call__bound__arguments = inspect.signature(_slugify_model_id).bind(model_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_slugify_model_id, 'tests.inference.unit_tests.test__slugify_model_id__unit_test_0', None, 'test_length_truncation_uses_original_for_digest', '_slugify_model_id', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_many_unique_short_model_ids_produce_unique_slugs_and_stable_digests():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1000
    model_ids = [f'model_{i}' for i in range(n)]
    results = [_slugify_model_id(mid) for mid in model_ids]
    for mid, out in zip(model_ids, results):
        expected_digest = hashlib.blake2s(mid.encode('utf-8'), digest_size=4).hexdigest()
        slug_portion = out[:-1 - len(expected_digest)]
        allowed_chars = set('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_-')
        if slug_portion != 'special-char-only-model-id':
            pass

def test_large_collection_including_long_and_special_inputs_no_exceptions_and_expected_format():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1000
    inputs = []
    for i in range(n):
        if i % 3 == 0:
            inputs.append('L' * (30 + i % 100))
        elif i % 3 == 1:
            inputs.append('!@#$$%^&*()_+' + str(i))
        else:
            inputs.append(f'mix-{i}_end')
    outputs = [_slugify_model_id(s) for s in inputs]
    for out in outputs:
        digest_part = out.rsplit('-', 1)[-1]
        int(digest_part, 16)