from __future__ import annotations

import builtins
import gc
import os
import time
import types

import pytest
from inference.core.workflows.errors import DynamicBlockError, PythonBlockError
from inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding import (
    _LOCAL_SHARED_GLOBALS, _get_python_code_imports, create_dynamic_module)
from inference.core.workflows.execution_engine.v1.dynamic_blocks.entities import \
    PythonCode


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 _set_mode(module, mode):
    """
    Helper to set the execution mode on the block_scaffolding module and
    return a function that restores the previous value.
    """
    prev = getattr(module, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', None)
    setattr(module, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', mode)

    def _restore():
        if prev is None:
            try:
                delattr(module, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE')
            except Exception:
                setattr(module, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', None)
        else:
            setattr(module, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', prev)
    return _restore

def test_basic_local_module_creation_and_run():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding as bs
    restore = _set_mode(bs, 'local')
    try:
        code = PythonCode(type='PythonCode', run_function_code='def run():\n    return 123\n', run_function_name='run', init_function_code=None, init_function_name='init', imports=[])
        codeflash_return_value = codeflash_wrap(create_dynamic_module, 'tests.inference.unit_tests.test_create_dynamic_module__unit_test_0', None, 'test_basic_local_module_creation_and_run', 'create_dynamic_module', '2_1', codeflash_loop_index, 'test_block', code, module_name='mod_test', api_key=None)
        codeflash_output = codeflash_return_value
        module = codeflash_output
    finally:
        restore()

def test_skip_class_eval_returns_stub_functions_and_init_stub():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding as bs
    restore = _set_mode(bs, 'local')
    try:
        code = PythonCode(type='PythonCode', run_function_code="def run():\n    return 'should not execute'\n", run_function_name='run', init_function_code="def init():\n    return {'x': 1}\n", init_function_name='init', imports=[])
        codeflash_return_value = codeflash_wrap(create_dynamic_module, 'tests.inference.unit_tests.test_create_dynamic_module__unit_test_0', None, 'test_skip_class_eval_returns_stub_functions_and_init_stub', 'create_dynamic_module', '2_1', codeflash_loop_index, 'test_block', code, module_name='stub_module', skip_class_eval=True)
        codeflash_output = codeflash_return_value
        module = codeflash_output
    finally:
        restore()

def test_syntax_error_at_exec_raises_pythonblockerror_with_lineno():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding as bs
    restore = _set_mode(bs, 'local')
    try:
        run_code = 'def run():\n    a =\n    return 1\n'
        code = PythonCode(type='PythonCode', run_function_code=run_code, run_function_name='run', init_function_code=None, init_function_name='init', imports=[])
        combined = _get_python_code_imports(code) + run_code
        try:
            compile(combined, '<string>', 'exec')
            pytest.fail('Expected compilation to fail due to syntax error, but it compiled fine.')
        except SyntaxError as se:
            expected_lineno = se.lineno
        with pytest.raises(PythonBlockError) as excinfo:
            codeflash_return_value = codeflash_wrap(create_dynamic_module, 'tests.inference.unit_tests.test_create_dynamic_module__unit_test_0', None, 'test_syntax_error_at_exec_raises_pythonblockerror_with_lineno', 'create_dynamic_module', '2_4_0', codeflash_loop_index, 'syntax_block', code, module_name='mod_syntax', api_key=None)
        err = excinfo.value
    finally:
        restore()

def test_exec_runtime_name_error_results_in_pythonblockerror_without_lineno_or_snippet():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding as bs
    restore = _set_mode(bs, 'local')
    try:
        run_code = 'THIS_IS_NOT_DEFINED = unknown_variable\n'
        code = PythonCode(type='PythonCode', run_function_code=run_code, run_function_name='run', init_function_code=None, init_function_name='init', imports=[])
        with pytest.raises(PythonBlockError) as excinfo:
            codeflash_return_value = codeflash_wrap(create_dynamic_module, 'tests.inference.unit_tests.test_create_dynamic_module__unit_test_0', None, 'test_exec_runtime_name_error_results_in_pythonblockerror_without_lineno_or_snippet', 'create_dynamic_module', '2_2_0', codeflash_loop_index, 'nameerror_block', code, module_name='mod_nameerror', api_key=None)
        err = excinfo.value
    finally:
        restore()

def test_large_scale_sum_1000_iterations_performs_correctly():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding as bs
    restore = _set_mode(bs, 'local')
    try:
        run_code = 'def run():\n    return sum(range(1000))\n'
        code = PythonCode(type='PythonCode', run_function_code=run_code, run_function_name='run', init_function_code=None, init_function_name='init', imports=[])
        codeflash_return_value = codeflash_wrap(create_dynamic_module, 'tests.inference.unit_tests.test_create_dynamic_module__unit_test_0', None, 'test_large_scale_sum_1000_iterations_performs_correctly', 'create_dynamic_module', '2_2', codeflash_loop_index, 'large_block', code, module_name='mod_large', api_key=None)
        codeflash_output = codeflash_return_value
        module = codeflash_output
    finally:
        restore()

def test_modal_mode_without_credentials_raises_dynamicblockerror():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding as bs
    restore = _set_mode(bs, 'modal')
    try:
        run_code = 'def run():\n    return 1\n'
        code = PythonCode(type='PythonCode', run_function_code=run_code, run_function_name='run', init_function_code=None, init_function_name='init', imports=[])
        with pytest.raises(DynamicBlockError):
            codeflash_return_value = codeflash_wrap(create_dynamic_module, 'tests.inference.unit_tests.test_create_dynamic_module__unit_test_0', None, 'test_modal_mode_without_credentials_raises_dynamicblockerror', 'create_dynamic_module', '2_2_0', codeflash_loop_index, 'modal_block', code, module_name='mod_modal', api_key=None)
    finally:
        restore()

def test_empty_run_function_code_results_in_no_run_attribute():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding as bs
    restore = _set_mode(bs, 'local')
    try:
        code = PythonCode(type='PythonCode', run_function_code='', run_function_name='run', init_function_code=None, init_function_name='init', imports=[])
        codeflash_return_value = codeflash_wrap(create_dynamic_module, 'tests.inference.unit_tests.test_create_dynamic_module__unit_test_0', None, 'test_empty_run_function_code_results_in_no_run_attribute', 'create_dynamic_module', '2_1', codeflash_loop_index, 'empty_block', code, module_name='mod_empty', api_key=None)
        codeflash_output = codeflash_return_value
        module = codeflash_output
    finally:
        restore()