from __future__ import annotations

import gc
import os
import time
from typing import Optional

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_code_with_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test basic traceback string generation with a single line of code.'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_basic_single_line_code_with_error', 'build_traceback_string', '1', codeflash_loop_index, code='x = 1 / 0', line_number=1, function_name='my_func', error_type='ZeroDivisionError', error_msg='division by zero')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_basic_multiple_lines_code_first_line_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test traceback with multiple lines of code, error on first line.'
    code = 'y = 2\nx = 1 / 0\nz = 3'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_basic_multiple_lines_code_first_line_error', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=1, function_name='test_func', error_type='ValueError', error_msg='invalid value')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_basic_multiple_lines_code_middle_line_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test traceback with multiple lines of code, error on middle line.'
    code = 'a = 1\nb = 2 + 3\nc = 4'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_basic_multiple_lines_code_middle_line_error', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=2, function_name='middle_func', error_type='TypeError', error_msg='type error occurred')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_basic_multiple_lines_code_last_line_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test traceback with multiple lines of code, error on last line.'
    code = 'x = 1\ny = 2\nz = undefined_var'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_basic_multiple_lines_code_last_line_error', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=3, function_name='last_line_func', error_type='NameError', error_msg="name 'undefined_var' is not defined")
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_code_line_is_stripped():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that extracted code line is stripped of leading/trailing whitespace.'
    code = '    x = 1 / 0    '
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_code_line_is_stripped', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=1, function_name='func', error_type='ZeroDivisionError', error_msg='division by zero')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_traceback_structure_and_order():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that the traceback structure follows the correct order.'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_traceback_structure_and_order', 'build_traceback_string', '1', codeflash_loop_index, code='error_line', line_number=1, function_name='func', error_type='RuntimeError', error_msg='test error')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    lines = result.split('\n')

def test_none_code():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with None as code input.'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_none_code', 'build_traceback_string', '1', codeflash_loop_index, code=None, line_number=1, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    lines = result.split('\n')

def test_empty_string_code():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with empty string as code input.'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_empty_string_code', 'build_traceback_string', '1', codeflash_loop_index, code='', line_number=1, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    lines = result.split('\n')

def test_line_number_out_of_bounds_too_high():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test when line_number is greater than the number of code lines.'
    code = 'line1\nline2'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_line_number_out_of_bounds_too_high', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=10, function_name='func', error_type='IndexError', error_msg='line out of bounds')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    lines = result.split('\n')

def test_line_number_zero():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with line_number of 0 (invalid/boundary case).'
    code = 'x = 1'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_line_number_zero', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=0, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    lines = result.split('\n')

def test_line_number_negative():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with negative line_number.'
    code = 'x = 1\ny = 2'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_line_number_negative', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=-1, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    lines = result.split('\n')

def test_empty_function_name():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with empty string as function name.'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_empty_function_name', 'build_traceback_string', '1', codeflash_loop_index, code='x = 1', line_number=1, function_name='', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_error_type():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with empty error type.'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_empty_error_type', 'build_traceback_string', '1', codeflash_loop_index, code='x = 1', line_number=1, function_name='func', error_type='', error_msg='error message')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_error_message():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with empty error message.'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_empty_error_message', 'build_traceback_string', '1', codeflash_loop_index, code='x = 1', line_number=1, function_name='func', error_type='RuntimeError', error_msg='')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_special_characters_in_code():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with special characters in code.'
    code = "print('hello\\nworld')"
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_special_characters_in_code', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=1, function_name='func', error_type='SyntaxError', error_msg='syntax error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_special_characters_in_error_message():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with special characters in error message.'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_special_characters_in_error_message', 'build_traceback_string', '1', codeflash_loop_index, code='x = 1', line_number=1, function_name='func', error_type='RuntimeError', error_msg='error with \'quotes\' and "double quotes"')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_unicode_characters_in_code():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with Unicode characters in code.'
    code = "message = '你好世界'"
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_unicode_characters_in_code', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=1, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_unicode_characters_in_error_message():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with Unicode characters in error message.'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_unicode_characters_in_error_message', 'build_traceback_string', '1', codeflash_loop_index, code='x = 1', line_number=1, function_name='func', error_type='RuntimeError', error_msg='错误消息')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_very_long_function_name():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with very long function name.'
    long_name = 'a' * 1000
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_very_long_function_name', 'build_traceback_string', '2', codeflash_loop_index, code='x = 1', line_number=1, function_name=long_name, error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_very_long_error_message():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with very long error message.'
    long_msg = 'error ' * 200
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_very_long_error_message', 'build_traceback_string', '2', codeflash_loop_index, code='x = 1', line_number=1, function_name='func', error_type='RuntimeError', error_msg=long_msg)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_code_with_only_whitespace():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with code that contains only whitespace.'
    code = '   \n  \n   '
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_code_with_only_whitespace', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=1, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    lines = result.split('\n')

def test_code_with_tabs_and_spaces():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with code containing mixed tabs and spaces.'
    code = '\t\tcode_with_tabs\t  '
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_code_with_tabs_and_spaces', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=1, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_multiline_code_with_empty_lines():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test multiline code with empty lines.'
    code = 'a = 1\n\nb = 2'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_multiline_code_with_empty_lines', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=2, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    lines = result.split('\n')

def test_function_name_with_special_chars():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test function name with special characters.'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_function_name_with_special_chars', 'build_traceback_string', '1', codeflash_loop_index, code='x = 1', line_number=1, function_name='func_<lambda>_test', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_error_type_with_brackets():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test error type that includes brackets or special formatting.'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_error_type_with_brackets', 'build_traceback_string', '1', codeflash_loop_index, code='x = 1', line_number=1, function_name='func', error_type='Error[CustomType]', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_very_large_code_single_error_line():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with very large code block but error on first line.'
    code_lines = [f'line{i} = {i}' for i in range(1000)]
    code = '\n'.join(code_lines)
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_very_large_code_single_error_line', 'build_traceback_string', '3', codeflash_loop_index, code=code, line_number=1, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_very_large_code_middle_error_line():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with very large code block, error on middle line.'
    code_lines = [f'line{i} = {i}' for i in range(1000)]
    code = '\n'.join(code_lines)
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_very_large_code_middle_error_line', 'build_traceback_string', '3', codeflash_loop_index, code=code, line_number=500, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_very_large_code_last_error_line():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with very large code block, error on last line.'
    code_lines = [f'line{i} = {i}' for i in range(1000)]
    code = '\n'.join(code_lines)
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_very_large_code_last_error_line', 'build_traceback_string', '3', codeflash_loop_index, code=code, line_number=1000, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_very_large_code_out_of_bounds():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with very large code block, line number out of bounds.'
    code_lines = [f'line{i} = {i}' for i in range(1000)]
    code = '\n'.join(code_lines)
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_very_large_code_out_of_bounds', 'build_traceback_string', '3', codeflash_loop_index, code=code, line_number=2000, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    lines = result.split('\n')

def test_code_with_very_long_lines():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with code containing very long individual lines.'
    long_line = 'x = ' + 'a' * 10000
    code = long_line
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_code_with_very_long_lines', 'build_traceback_string', '3', codeflash_loop_index, code=code, line_number=1, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_multiple_very_large_code_blocks_sequential():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test multiple calls with large code blocks in sequence.'
    code = '\n'.join([f'line{i} = {i}' for i in range(500)])
    for line_num in [1, 100, 250, 500]:
        codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_multiple_very_large_code_blocks_sequential', 'build_traceback_string', '2_0', codeflash_loop_index, code=code, line_number=line_num, function_name='func', error_type='RuntimeError', error_msg='error')
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_code_with_many_unicode_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with code containing many Unicode characters.'
    code = '\n'.join([f"line{i} = '中文{i}'" for i in range(500)])
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_code_with_many_unicode_characters', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=250, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_extremely_large_error_message():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with extremely large error message (100KB).'
    large_msg = 'error ' * 10000
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_extremely_large_error_message', 'build_traceback_string', '2', codeflash_loop_index, code='x = 1', line_number=1, function_name='func', error_type='RuntimeError', error_msg=large_msg)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_nested_multiline_code_structure():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with code that has complex nested multiline structure.'
    code = 'def outer():\n    def inner():\n        try:\n            x = 1 / 0\n        except ZeroDivisionError:\n            pass\n        finally:\n            y = 2\n    return inner()\n\nclass MyClass:\n    def method(self):\n        for i in range(10):\n            if i > 5:\n                z = [1, 2, 3]'
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_nested_multiline_code_structure', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=4, function_name='inner', error_type='ZeroDivisionError', error_msg='division by zero')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_code_with_all_special_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with code containing a wide variety of special characters.'
    code = "x = [1, 2, 3]; y = {'a': 1}; z = lambda x: x + 1; msg = r'C:\\path\\to\\file'"
    codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_code_with_all_special_characters', 'build_traceback_string', '2', codeflash_loop_index, code=code, line_number=1, function_name='func', error_type='RuntimeError', error_msg='error')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_stress_many_function_calls():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Stress test with many sequential function calls.'
    code = 'line1\nline2\nline3'
    for i in range(1000):
        line_num = i % 3 + 1
        codeflash_return_value = codeflash_wrap(build_traceback_string, 'tests.inference.unit_tests.test_build_traceback_string__unit_test_1', None, 'test_stress_many_function_calls', 'build_traceback_string', '2_1', codeflash_loop_index, code=code, line_number=line_num, function_name=f'func{i}', error_type=f'Error{i}', error_msg=f'error {i}')
        codeflash_output = codeflash_return_value
        result = codeflash_output