from __future__ import annotations

import gc
import os
import time
from typing import Any, Dict, List

import pytest
from inference.core.workflows.core_steps.common.query_language.entities.introspection import \
    OperationDescription
from inference.core.workflows.core_steps.common.query_language.entities.operations import \
    OperationsChain
from inference.core.workflows.core_steps.common.query_language.introspection.core import \
    prepare_operations_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}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

def _make_type_def(type_name: str, *, property_name_schema: Dict[str, Any]=None, compound: bool=False, input_kind: List=None, output_kind: List=None, nested_in: List=None, nested_out: List=None) -> Dict[str, Any]:
    """
    Build a minimal type definition dict that matches what prepare_operations_descriptions
    expects from a $defs entry. All lists default to empty lists (which are acceptable).
    """
    input_kind = input_kind if input_kind is not None else []
    output_kind = output_kind if output_kind is not None else []
    props = {'type': {'const': type_name}}
    if property_name_schema is not None:
        props['property_name'] = property_name_schema
    td = {'properties': props, 'input_kind': input_kind, 'output_kind': output_kind}
    if compound:
        td['compound'] = True
    if nested_in is not None:
        td['nested_operation_input_kind'] = nested_in
    if nested_out is not None:
        td['nested_operation_output_kind'] = nested_out
    return td

def test_single_operation_minimal_schema_property_name_none(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    defs = {'OpA': _make_type_def('opA')}
    schema = {'properties': {'operations': {'items': {'oneOf': [{'$ref': '#/$defs/OpA'}]}}}, '$defs': defs}
    monkeypatch.setattr(OperationsChain, 'model_json_schema', classmethod(lambda cls: schema), raising=True)
    codeflash_return_value = codeflash_wrap(prepare_operations_descriptions, 'tests.inference.unit_tests.test_prepare_operations_descriptions__unit_test_0', None, 'test_single_operation_minimal_schema_property_name_none', 'prepare_operations_descriptions', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    desc = result[0]

def test_property_name_enum_resolves_from_defs_ref(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    defs = {'OperationX': _make_type_def('opx', property_name_schema={'$ref': '#/$defs/PropertyNames'}, compound=True), 'PropertyNames': {'enum': ['size', 'color', 'weight']}}
    schema = {'properties': {'operations': {'items': {'oneOf': [{'$ref': '#/$defs/OperationX'}]}}}, '$defs': defs}
    monkeypatch.setattr(OperationsChain, 'model_json_schema', classmethod(lambda cls: schema), raising=True)
    codeflash_return_value = codeflash_wrap(prepare_operations_descriptions, 'tests.inference.unit_tests.test_prepare_operations_descriptions__unit_test_0', None, 'test_property_name_enum_resolves_from_defs_ref', 'prepare_operations_descriptions', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    desc = result[0]

def test_property_name_enum_resolves_from_direct_enum_in_schema(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    defs = {'OpDirectEnum': _make_type_def('op_direct', property_name_schema={'enum': ['height', 'width']})}
    schema = {'properties': {'operations': {'items': {'oneOf': [{'$ref': '#/$defs/OpDirectEnum'}]}}}, '$defs': defs}
    monkeypatch.setattr(OperationsChain, 'model_json_schema', classmethod(lambda cls: schema), raising=True)
    codeflash_return_value = codeflash_wrap(prepare_operations_descriptions, 'tests.inference.unit_tests.test_prepare_operations_descriptions__unit_test_0', None, 'test_property_name_enum_resolves_from_direct_enum_in_schema', 'prepare_operations_descriptions', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    desc = result[0]

def test_no_operations_oneOf_empty_returns_empty_list(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'properties': {'operations': {'items': {'oneOf': []}}}, '$defs': {}}
    monkeypatch.setattr(OperationsChain, 'model_json_schema', classmethod(lambda cls: schema), raising=True)
    codeflash_return_value = codeflash_wrap(prepare_operations_descriptions, 'tests.inference.unit_tests.test_prepare_operations_descriptions__unit_test_0', None, 'test_no_operations_oneOf_empty_returns_empty_list', 'prepare_operations_descriptions', '2', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_ignores_oneOf_items_without_ref(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    defs = {'OpGood': _make_type_def('op_good')}
    schema = {'properties': {'operations': {'items': {'oneOf': [{'$ref': '#/$defs/OpGood'}, {'some': 'structure'}, {}]}}}, '$defs': defs}
    monkeypatch.setattr(OperationsChain, 'model_json_schema', classmethod(lambda cls: schema), raising=True)
    codeflash_return_value = codeflash_wrap(prepare_operations_descriptions, 'tests.inference.unit_tests.test_prepare_operations_descriptions__unit_test_0', None, 'test_ignores_oneOf_items_without_ref', 'prepare_operations_descriptions', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_number_of_operations_scaling(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1000
    defs: Dict[str, Any] = {}
    one_of_list: List[Dict[str, str]] = []
    for i in range(n):
        name = f'Op{i}'
        op_const = f'op_{i}'
        defs[name] = _make_type_def(op_const)
        one_of_list.append({'$ref': f'#/$defs/{name}'})
    schema = {'properties': {'operations': {'items': {'oneOf': one_of_list}}}, '$defs': defs}
    monkeypatch.setattr(OperationsChain, 'model_json_schema', classmethod(lambda cls: schema), raising=True)
    codeflash_return_value = codeflash_wrap(prepare_operations_descriptions, 'tests.inference.unit_tests.test_prepare_operations_descriptions__unit_test_0', None, 'test_large_number_of_operations_scaling', 'prepare_operations_descriptions', '6', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    returned_types = {d.operation_type for d in result}
    expected_types = {f'op_{i}' for i in range(n)}