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 ExecutionEngineRuntimeError
from inference.core.workflows.execution_engine.entities.base import (
    CoordinatesSystem, JsonField)
from inference.core.workflows.execution_engine.profiling.core import \
    NullWorkflowsProfiler
from inference.core.workflows.execution_engine.v1.compiler.entities import \
    CompiledWorkflow
from inference.core.workflows.execution_engine.v1.executor.core import \
    _run_workflow
from networkx import DiGraph


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 MinimalWorkflowDefinition:

    def __init__(self, outputs):
        self.outputs = outputs

def make_compiled_workflow_with_outputs(outputs):
    """
    Helper to create a CompiledWorkflow instance with:
      - a simple execution_graph (empty DiGraph suffices for tests that do not
        exercise step execution),
      - a minimal workflow_definition object that exposes .outputs,
      - empty steps, substitutions and other metadata expected by CompiledWorkflow.
    """
    execution_graph = DiGraph()
    workflow_definition = MinimalWorkflowDefinition(outputs=outputs)
    return CompiledWorkflow(workflow_definition=workflow_definition, execution_graph=execution_graph, steps={}, input_substitutions=[], workflow_json={}, init_parameters={}, kinds_serializers={}, kinds_deserializers={})

def test_run_workflow_returns_single_empty_result_for_no_outputs():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    workflow = make_compiled_workflow_with_outputs(outputs=[])
    runtime_parameters = {}
    profiler = NullWorkflowsProfiler.init()
    _call__bound__arguments = inspect.signature(_run_workflow).bind(workflow=workflow, runtime_parameters=runtime_parameters, max_concurrent_steps=1, kinds_serializers=None, serialize_results=False, profiler=profiler, executor=None, step_error_handler=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_run_workflow, 'tests.inference.unit_tests.test__run_workflow__unit_test_0', None, 'test_run_workflow_returns_single_empty_result_for_no_outputs', '_run_workflow', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    results = codeflash_output

def test_run_workflow_raises_on_invalid_selector_in_output():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    invalid_output = JsonField(type='JsonField', name='bad', selector='not_a_selector')
    workflow = make_compiled_workflow_with_outputs(outputs=[invalid_output])
    profiler = NullWorkflowsProfiler.init()
    with pytest.raises(ExecutionEngineRuntimeError):
        _call__bound__arguments = inspect.signature(_run_workflow).bind(workflow=workflow, runtime_parameters={}, max_concurrent_steps=1, kinds_serializers=None, serialize_results=False, profiler=profiler, executor=None, step_error_handler=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(_run_workflow, 'tests.inference.unit_tests.test__run_workflow__unit_test_0', None, 'test_run_workflow_raises_on_invalid_selector_in_output', '_run_workflow', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_run_workflow_many_empty_runs_is_deterministic_and_fast():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    workflow = make_compiled_workflow_with_outputs(outputs=[])
    profiler = NullWorkflowsProfiler.init()
    for _ in range(100):
        _call__bound__arguments = inspect.signature(_run_workflow).bind(workflow=workflow, runtime_parameters={}, max_concurrent_steps=2, kinds_serializers=None, serialize_results=False, profiler=profiler, executor=None, step_error_handler=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(_run_workflow, 'tests.inference.unit_tests.test__run_workflow__unit_test_0', None, 'test_run_workflow_many_empty_runs_is_deterministic_and_fast', '_run_workflow', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output