from __future__ import annotations

import gc
import os
import time

import pytest
from inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding import \
    build_traceback_string


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_basic_single_line_traceback():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    code = 'x = 1\ny = 2\nz = x + y'
    line_number = 3
    function_name = 'my_func'
    error_type = 'ValueError'
    error_msg = 'boom'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_0', None, 'test_basic_single_line_traceback', 'build_traceback_string', '5', codeflash_loop_index, code, line_number, function_name, error_type, error_msg)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = 'Traceback (most recent call last):\n  File "Python Block", line 3, in my_func\n    z = x + y\nValueError: boom'

def test_no_code_line_when_line_out_of_range_or_zero():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    code = 'a = 10\nb = 20'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_0', None, 'test_no_code_line_when_line_out_of_range_or_zero', 'build_traceback_string', '1', codeflash_loop_index, code, 0, 'fn', 'TypeError', 't')
    codeflash_output = codeflash_return_value
    result_zero = codeflash_output
    expected_zero = 'Traceback (most recent call last):\n  File "Python Block", line 0, in fn\nTypeError: t'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_0', None, 'test_no_code_line_when_line_out_of_range_or_zero', 'build_traceback_string', '4', codeflash_loop_index, code, 10, 'fn', 'TypeError', 't')
    codeflash_output = codeflash_return_value
    result_large = codeflash_output
    expected_large = 'Traceback (most recent call last):\n  File "Python Block", line 10, in fn\nTypeError: t'

def test_none_code_behaves_like_empty_code():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_0', None, 'test_none_code_behaves_like_empty_code', 'build_traceback_string', '0', codeflash_loop_index, None, 1, 'name', 'Err', 'msg')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = 'Traceback (most recent call last):\n  File "Python Block", line 1, in name\nErr: msg'

def test_whitespace_only_target_line_is_omitted():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    code = 'first_line\n   \t   \nthird_line'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_0', None, 'test_whitespace_only_target_line_is_omitted', 'build_traceback_string', '1', codeflash_loop_index, code, 2, 'whitespace_fn', 'RuntimeError', 'oops')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = 'Traceback (most recent call last):\n  File "Python Block", line 2, in whitespace_fn\nRuntimeError: oops'

def test_unicode_and_special_characters_preserved():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    code = "print('✓')\n# a comment with emoji 😄"
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_0', None, 'test_unicode_and_special_characters_preserved', 'build_traceback_string', '1', codeflash_loop_index, code, 1, 'функция', 'CustomError', 'провал 🌶️')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = 'Traceback (most recent call last):\n  File "Python Block", line 1, in функция\n    print(\'✓\')\nCustomError: провал 🌶️'

def test_empty_error_type_or_message_still_formats_colon():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    code = 'x = 2'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_0', None, 'test_empty_error_type_or_message_still_formats_colon', 'build_traceback_string', '1', codeflash_loop_index, code, 1, 'fn', '', 'only message')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = 'Traceback (most recent call last):\n  File "Python Block", line 1, in fn\n    x = 2\n: only message'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_0', None, 'test_empty_error_type_or_message_still_formats_colon', 'build_traceback_string', '4', codeflash_loop_index, code, 1, 'fn', 'ErrType', '')
    codeflash_output = codeflash_return_value
    result2 = codeflash_output
    expected2 = 'Traceback (most recent call last):\n  File "Python Block", line 1, in fn\n    x = 2\nErrType: '

def test_negative_line_number_treated_as_out_of_range():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    code = 'only_line'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_0', None, 'test_negative_line_number_treated_as_out_of_range', 'build_traceback_string', '1', codeflash_loop_index, code, -5, 'neg', 'Err', 'neg')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = 'Traceback (most recent call last):\n  File "Python Block", line -5, in neg\nErr: neg'

def test_large_code_block_with_many_lines_and_many_calls():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    lines = [f'line_{i}' for i in range(1, 1001)]
    large_code = '\n'.join(lines)
    for i in range(1, 1001):
        codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_0', None, 'test_large_code_block_with_many_lines_and_many_calls', 'build_traceback_string', '2_0', codeflash_loop_index, large_code, i, f'fn_{i}', 'LargeError', f'msg_{i}')
        codeflash_output = codeflash_return_value
        result = codeflash_output
        expected_prefix = f'Traceback (most recent call last):\n  File "Python Block", line {i}, in fn_{i}\n    line_{i}\nLargeError: msg_{i}'