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.core_steps.common.query_language.entities import \
    operations as ops_module
from inference.core.workflows.core_steps.common.query_language.entities.introspection import \
    OperatorDescription
from inference.core.workflows.core_steps.common.query_language.entities.operations import (
    DoesNotExist, Equals, Exists, In, IsEmpty, IsFalse, IsNotEmpty, IsTrue,
    NotEquals, NumberGreater, NumberGreaterEqual, NumberLower,
    NumberLowerEqual, StringContains, StringEndsWith, StringStartsWith)
from inference.core.workflows.core_steps.common.query_language.introspection.core import \
    prepare_operators_descriptions


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
_OPERATOR_CLASSES_IN_ORDER = [In, StringContains, StringEndsWith, StringStartsWith, NumberLowerEqual, NumberLower, NumberGreaterEqual, NumberGreater, NotEquals, Equals, Exists, DoesNotExist, IsTrue, IsFalse, IsEmpty, IsNotEmpty]

def _collect_aliases_in_order():
    """
    Helper that reads each operator class's model_json_schema and collects the 'type'
    enum aliases in the same nested loop order as prepare_operators_descriptions:
    iterate operator classes in the same order, then each alias in the enum order.
    """
    aliases = []
    for cls in _OPERATOR_CLASSES_IN_ORDER:
        schema = cls.model_json_schema()
        enum_aliases = schema['properties']['type'].get('enum', [])
        aliases.extend(enum_aliases)
    return aliases

def test_basic_returns_list_of_operator_descriptions():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(prepare_operators_descriptions).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(prepare_operators_descriptions, 'tests.inference.unit_tests.test_prepare_operators_descriptions__unit_test_0', None, 'test_basic_returns_list_of_operator_descriptions', 'prepare_operators_descriptions', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    results = codeflash_output
    expected_aliases = _collect_aliases_in_order()
    returned_aliases = [item.operator_type for item in results]

def test_operator_description_fields_match_schema_for_each_alias():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(prepare_operators_descriptions).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(prepare_operators_descriptions, 'tests.inference.unit_tests.test_prepare_operators_descriptions__unit_test_0', None, 'test_operator_description_fields_match_schema_for_each_alias', 'prepare_operators_descriptions', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    results = codeflash_output
    alias_to_schema = {}
    for cls in _OPERATOR_CLASSES_IN_ORDER:
        schema = cls.model_json_schema()
        enum_aliases = schema['properties']['type'].get('enum', [])
        for alias in enum_aliases:
            alias_to_schema[alias] = schema
    for desc in results:
        schema = alias_to_schema[desc.operator_type]
        expected_description = schema.get('description')

def test_multiple_aliases_per_operator_are_all_included():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(prepare_operators_descriptions).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(prepare_operators_descriptions, 'tests.inference.unit_tests.test_prepare_operators_descriptions__unit_test_0', None, 'test_multiple_aliases_per_operator_are_all_included', 'prepare_operators_descriptions', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    results = codeflash_output
    returned_aliases = {r.operator_type for r in results}
    equals_aliases = Equals.model_json_schema()['properties']['type'].get('enum', [])
    not_equals_aliases = NotEquals.model_json_schema()['properties']['type'].get('enum', [])
    for alias in equals_aliases:
        pass
    for alias in not_equals_aliases:
        pass

def test_ordering_matches_operator_types_and_enum_order():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(prepare_operators_descriptions).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(prepare_operators_descriptions, 'tests.inference.unit_tests.test_prepare_operators_descriptions__unit_test_0', None, 'test_ordering_matches_operator_types_and_enum_order', 'prepare_operators_descriptions', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    results = codeflash_output
    returned_aliases = [r.operator_type for r in results]
    expected_aliases_in_order = _collect_aliases_in_order()

def test_idempotent_calls_and_fresh_instances():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(prepare_operators_descriptions).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(prepare_operators_descriptions, 'tests.inference.unit_tests.test_prepare_operators_descriptions__unit_test_0', None, 'test_idempotent_calls_and_fresh_instances', 'prepare_operators_descriptions', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    first = codeflash_output
    _call__bound__arguments = inspect.signature(prepare_operators_descriptions).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(prepare_operators_descriptions, 'tests.inference.unit_tests.test_prepare_operators_descriptions__unit_test_0', None, 'test_idempotent_calls_and_fresh_instances', 'prepare_operators_descriptions', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    second = codeflash_output
    first_types = [o.operator_type for o in first]
    second_types = [o.operator_type for o in second]
    for a, b in zip(first, second):
        pass

def test_operands_number_values_sensible_and_consistent():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(prepare_operators_descriptions).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(prepare_operators_descriptions, 'tests.inference.unit_tests.test_prepare_operators_descriptions__unit_test_0', None, 'test_operands_number_values_sensible_and_consistent', 'prepare_operators_descriptions', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    results = codeflash_output
    for desc in results:
        pass

def test_large_scale_repeated_invocations_consistent():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(prepare_operators_descriptions).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(prepare_operators_descriptions, 'tests.inference.unit_tests.test_prepare_operators_descriptions__unit_test_0', None, 'test_large_scale_repeated_invocations_consistent', 'prepare_operators_descriptions', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    baseline = codeflash_output
    baseline_types = [o.operator_type for o in baseline]
    for _ in range(1000):
        _call__bound__arguments = inspect.signature(prepare_operators_descriptions).bind()
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(prepare_operators_descriptions, 'tests.inference.unit_tests.test_prepare_operators_descriptions__unit_test_0', None, 'test_large_scale_repeated_invocations_consistent', 'prepare_operators_descriptions', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        current = codeflash_output
        current_types = [o.operator_type for o in current]