from __future__ import annotations

import gc
import os
import time
from types import ModuleType

import pytest
from inference.core.telemetry import _ForceTraceRootSampler


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 _make_delegate_with_callable(fn):
    """
    Helper: construct a ModuleType instance and attach a get_description
    callable attribute. We use ModuleType because it's a real stdlib class
    that allows setting attributes (unlike built-in object()) and is not a
    test stub class defined by us.
    """
    m = ModuleType('delegate_module')
    m.get_description = fn
    return m

def test_basic_description_concatenation():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    delegate = _make_delegate_with_callable(lambda: 'delegate')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_basic_description_concatenation', 'get_description', '2', codeflash_loop_index)
    codeflash_output = codeflash_return_value

def test_empty_delegate_description():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    delegate = _make_delegate_with_callable(lambda: '')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_empty_delegate_description', 'get_description', '2', codeflash_loop_index)
    codeflash_output = codeflash_return_value

def test_numeric_delegate_description_converted_to_string():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    delegate = _make_delegate_with_callable(lambda: 123)
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_numeric_delegate_description_converted_to_string', 'get_description', '2', codeflash_loop_index)
    codeflash_output = codeflash_return_value

def test_none_delegate_description_converted_to_string():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    delegate = _make_delegate_with_callable(lambda: None)
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_none_delegate_description_converted_to_string', 'get_description', '2', codeflash_loop_index)
    codeflash_output = codeflash_return_value

def test_special_characters_in_delegate_description():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    special = 'weird_chars: ) ( ☃ , \n \t'
    delegate = _make_delegate_with_callable(lambda: special)
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_special_characters_in_delegate_description', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value

def test_delegate_raises_propagates_exception():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

    def raiser():
        raise ValueError('boom')
    delegate = _make_delegate_with_callable(raiser)
    sampler = _ForceTraceRootSampler(delegate)
    with pytest.raises(ValueError):
        codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_delegate_raises_propagates_exception', 'get_description', '3_0', codeflash_loop_index)

def test_nested_samplers_produce_nested_descriptions():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inner_delegate = _make_delegate_with_callable(lambda: 'inner')
    inner_sampler = _ForceTraceRootSampler(inner_delegate)
    outer_sampler = _ForceTraceRootSampler(inner_sampler)
    codeflash_return_value = codeflash_wrap(outer_sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_nested_samplers_produce_nested_descriptions', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value

def test_idempotent_multiple_calls_return_same_string():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    delegate = _make_delegate_with_callable(lambda: 'stable')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_idempotent_multiple_calls_return_same_string', 'get_description', '2', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    first = codeflash_output
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_idempotent_multiple_calls_return_same_string', 'get_description', '4', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    second = codeflash_output
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_idempotent_multiple_calls_return_same_string', 'get_description', '6', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    third = codeflash_output

def test_delegate_returns_long_string():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    long_desc = 'x' * 10000
    delegate = _make_delegate_with_callable(lambda: long_desc)
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_delegate_returns_long_string', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value

def test_delegate_returns_non_str_iterable():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    delegate = _make_delegate_with_callable(lambda: ('a', 1, None))
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_delegate_returns_non_str_iterable', 'get_description', '2', codeflash_loop_index)
    codeflash_output = codeflash_return_value

def test_delegate_is_another_module_with_stateful_callable():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    state = {'count': 0}

    def stateful():
        state['count'] += 1
        return f"count={state['count']}"
    delegate = _make_delegate_with_callable(stateful)
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_delegate_is_another_module_with_stateful_callable', 'get_description', '4', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_delegate_is_another_module_with_stateful_callable', 'get_description', '5', codeflash_loop_index)
    codeflash_output = codeflash_return_value

def test_large_scale_many_samplers_consistent_small_descriptions():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1000
    samplers = []
    for i in range(n):
        delegate = _make_delegate_with_callable(lambda v=i: f'd{v}')
        samplers.append(_ForceTraceRootSampler(delegate))
    for i, sampler in enumerate(samplers):
        codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_large_scale_many_samplers_consistent_small_descriptions', 'get_description', '3_0', codeflash_loop_index)
        codeflash_output = codeflash_return_value

def test_large_scale_long_descriptions_mixture():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 200
    samplers = []
    for i in range(1, n + 1):
        long_part = 'x' * i
        delegate = _make_delegate_with_callable(lambda s=long_part: s)
        samplers.append(_ForceTraceRootSampler(delegate))
    for i, sampler in enumerate(samplers, start=1):
        expected = 'ForceTraceRootSampler(' + 'x' * i + ')'
        codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_0', None, 'test_large_scale_long_descriptions_mixture', 'get_description', '3_1', codeflash_loop_index)
        codeflash_output = codeflash_return_value