from __future__ import annotations

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

import dill as pickle
import pytest
from inference.core.workflows.errors import (
    DynamicBlockError, PythonBlockError, WorkflowEnvironmentConfigurationError)
from inference.core.workflows.execution_engine.v1.dynamic_blocks import \
    block_scaffolding as bs
from inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding import \
    assembly_custom_python_block
from inference.core.workflows.execution_engine.v1.dynamic_blocks.entities import \
    PythonCode
from inference.core.workflows.prototypes.block import WorkflowBlockManifest


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

class MinimalManifest(WorkflowBlockManifest):
    type: str = 'test_manifest'
    name: str = 'minimal'

    @classmethod
    def describe_outputs(cls):
        return []

def test_basic_run_returns_expected_value(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(bs, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', 'local')
    monkeypatch.setattr(bs, 'ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS', True)
    python_code = PythonCode(type='PythonCode', imports=[], run_function_code='\ndef run(self, x):\n    # simple computation to verify call semantics\n    return {"res": x + 1}\n', run_function_name='run', init_function_code=None, init_function_name='init')
    _call__bound__arguments = inspect.signature(bs.assembly_custom_python_block).bind(block_type_name='increment_block', unique_identifier='uid_basic', manifest=MinimalManifest, python_code=python_code)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(bs.assembly_custom_python_block, 'tests.inference.unit_tests.test_assembly_custom_python_block__unit_test_0', None, 'test_basic_run_returns_expected_value', 'assembly_custom_python_block', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    Dynamic = codeflash_output
    instance = Dynamic(api_key=None)
    result = instance.run(x=10)

def test_get_init_parameters_and_get_manifest(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(bs, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', 'local')
    monkeypatch.setattr(bs, 'ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS', True)
    python_code = PythonCode(type='PythonCode', imports=[], run_function_code='\ndef run(self):\n    return {"ok": True}\n', run_function_name='run', init_function_code=None, init_function_name='init')
    _call__bound__arguments = inspect.signature(bs.assembly_custom_python_block).bind(block_type_name='meta_block', unique_identifier='uid_meta', manifest=MinimalManifest, python_code=python_code)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(bs.assembly_custom_python_block, 'tests.inference.unit_tests.test_assembly_custom_python_block__unit_test_0', None, 'test_get_init_parameters_and_get_manifest', 'assembly_custom_python_block', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    Dynamic = codeflash_output

def test_run_raises_when_local_execution_disallowed(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(bs, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', 'local')
    monkeypatch.setattr(bs, 'ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS', False)
    python_code = PythonCode(type='PythonCode', imports=[], run_function_code='\ndef run(self):\n    return {"should_not": "run"}\n', run_function_name='run', init_function_code=None, init_function_name='init')
    _call__bound__arguments = inspect.signature(bs.assembly_custom_python_block).bind(block_type_name='blocked_block', unique_identifier='uid_blocked', manifest=MinimalManifest, python_code=python_code)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(bs.assembly_custom_python_block, 'tests.inference.unit_tests.test_assembly_custom_python_block__unit_test_0', None, 'test_run_raises_when_local_execution_disallowed', 'assembly_custom_python_block', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    Dynamic = codeflash_output
    inst = Dynamic(api_key=None)
    with pytest.raises(WorkflowEnvironmentConfigurationError):
        inst.run()

def test_missing_declared_run_function_raises_dynamic_block_error(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(bs, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', 'local')
    monkeypatch.setattr(bs, 'ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS', True)
    python_code = PythonCode(type='PythonCode', imports=[], run_function_code='\ndef run(self):\n    return {"ok": True}\n', run_function_name='nonexistent', init_function_code=None, init_function_name='init')
    with pytest.raises(DynamicBlockError):
        _call__bound__arguments = inspect.signature(bs.assembly_custom_python_block).bind(block_type_name='bad_declared_fn', unique_identifier='uid_missing', manifest=MinimalManifest, python_code=python_code)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(bs.assembly_custom_python_block, 'tests.inference.unit_tests.test_assembly_custom_python_block__unit_test_0', None, 'test_missing_declared_run_function_raises_dynamic_block_error', 'assembly_custom_python_block', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_syntax_error_in_code_raises_python_block_error(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(bs, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', 'local')
    monkeypatch.setattr(bs, 'ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS', True)
    python_code = PythonCode(type='PythonCode', imports=[], run_function_code='\ndef run(self):\n    x =\n    return 1\n', run_function_name='run', init_function_code=None, init_function_name='init')
    with pytest.raises(PythonBlockError):
        _call__bound__arguments = inspect.signature(bs.assembly_custom_python_block).bind(block_type_name='syntax_error_block', unique_identifier='uid_syntax', manifest=MinimalManifest, python_code=python_code)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(bs.assembly_custom_python_block, 'tests.inference.unit_tests.test_assembly_custom_python_block__unit_test_0', None, 'test_syntax_error_in_code_raises_python_block_error', 'assembly_custom_python_block', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_modal_mode_without_modal_credentials_raises_dynamic_block_error(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(bs, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', 'modal')
    python_code = PythonCode(type='PythonCode', imports=[], run_function_code='\ndef run(self):\n    return {"ok": True}\n', run_function_name='run', init_function_code=None, init_function_name='init')
    with pytest.raises(DynamicBlockError):
        _call__bound__arguments = inspect.signature(bs.assembly_custom_python_block).bind(block_type_name='modal_block', unique_identifier='uid_modal', manifest=MinimalManifest, python_code=python_code, api_key=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(bs.assembly_custom_python_block, 'tests.inference.unit_tests.test_assembly_custom_python_block__unit_test_0', None, 'test_modal_mode_without_modal_credentials_raises_dynamic_block_error', 'assembly_custom_python_block', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_none_python_code_raises_attribute_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    with pytest.raises(AttributeError):
        _call__bound__arguments = inspect.signature(bs.assembly_custom_python_block).bind(block_type_name='none_block', unique_identifier='uid_none', manifest=MinimalManifest, python_code=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(bs.assembly_custom_python_block, 'tests.inference.unit_tests.test_assembly_custom_python_block__unit_test_0', None, 'test_none_python_code_raises_attribute_error', 'assembly_custom_python_block', '0_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_runtime_exception_from_user_code_is_wrapped_and_contains_stdout_and_stderr(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(bs, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', 'local')
    monkeypatch.setattr(bs, 'ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS', True)
    python_code = PythonCode(type='PythonCode', imports=[], run_function_code='\ndef run(self):\n    # Write to stdout\n    print("hello_stdout")\n    # Write to stderr\n    import sys\n    sys.stderr.write("hello_stderr\\n")\n    # Then raise an exception to trigger error handling path\n    raise ValueError("user_error")\n', run_function_name='run', init_function_code=None, init_function_name='init')
    _call__bound__arguments = inspect.signature(bs.assembly_custom_python_block).bind(block_type_name='error_out_block', unique_identifier='uid_err', manifest=MinimalManifest, python_code=python_code)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(bs.assembly_custom_python_block, 'tests.inference.unit_tests.test_assembly_custom_python_block__unit_test_0', None, 'test_runtime_exception_from_user_code_is_wrapped_and_contains_stdout_and_stderr', 'assembly_custom_python_block', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    Dynamic = codeflash_output
    inst = Dynamic(api_key=None)
    with pytest.raises(PythonBlockError) as exc_info:
        inst.run()
    err = exc_info.value

def test_large_scale_processing_sum_of_1000(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(bs, 'WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE', 'local')
    monkeypatch.setattr(bs, 'ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS', True)
    python_code = PythonCode(type='PythonCode', imports=[], run_function_code='\ndef run(self, items):\n    # A simple loop over items to emulate realistic processing\n    total = 0\n    for v in items:\n        total += v\n    return {"sum": total, "count": len(items)}\n', run_function_name='run', init_function_code=None, init_function_name='init')
    _call__bound__arguments = inspect.signature(bs.assembly_custom_python_block).bind(block_type_name='sum_block', unique_identifier='uid_large', manifest=MinimalManifest, python_code=python_code)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(bs.assembly_custom_python_block, 'tests.inference.unit_tests.test_assembly_custom_python_block__unit_test_0', None, 'test_large_scale_processing_sum_of_1000', 'assembly_custom_python_block', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    Dynamic = codeflash_output
    inst = Dynamic(api_key=None)
    items = list(range(1000))
    result = inst.run(items=items)