from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path
from typing import List, Set, Tuple

import dill as pickle
import pytest
from inference.core.workflows.execution_engine.entities.base import Batch


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

def test_remove_no_indices_returns_identical_batch():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    content = [10, 20, 30]
    indices: List[Tuple[int, ...]] = [(0,), (1,), (2,)]
    batch = Batch(content=content[:], indices=indices[:])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=set())
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_0', None, 'test_remove_no_indices_returns_identical_batch', 'remove_by_indices', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_nonexistent_indices_does_nothing():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    content = ['a', 'b']
    indices = [(0,), (1,)]
    batch = Batch(content=content[:], indices=indices[:])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(99,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_0', None, 'test_remove_nonexistent_indices_does_nothing', 'remove_by_indices', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_all_indices_returns_empty_batch():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    content = [1, 2, 3, 4]
    indices = [(0,), (1,), (2,), (3,)]
    batch = Batch(content=content[:], indices=indices[:])
    removal_set: Set[Tuple[int, ...]] = set(indices)
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=removal_set)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_0', None, 'test_remove_all_indices_returns_empty_batch', 'remove_by_indices', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_indices_none_raises_type_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    batch = Batch(content=[1, 2, 3], indices=None)
    with pytest.raises(TypeError):
        _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=set())
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_0', None, 'test_indices_none_raises_type_error', 'remove_by_indices', '1_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_nested_batch_removes_inner_element_only_when_matching_full_tuple():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    outer_content = ['outer0']
    nested_content = [100, 200]
    nested_indices = [(1, 0), (1, 1)]
    nested_batch = Batch(content=nested_content[:], indices=nested_indices[:])
    outer_content.append(nested_batch)
    outer_content.append('outer2')
    outer_indices = [(0,), (1,), (2,)]
    outer = Batch(content=outer_content[:], indices=outer_indices[:])
    _call__bound__arguments = inspect.signature(outer.remove_by_indices).bind(indices_to_remove={(1, 1)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(outer.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_0', None, 'test_nested_batch_removes_inner_element_only_when_matching_full_tuple', 'remove_by_indices', '8', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    kept_nested = result._content[1]

def test_nested_batch_removes_entire_nested_when_parent_index_removed():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    nested = Batch(content=[10, 20], indices=[(1, 0), (1, 1)])
    outer = Batch(content=['x', nested, 'z'], indices=[(0,), (1,), (2,)])
    _call__bound__arguments = inspect.signature(outer.remove_by_indices).bind(indices_to_remove={(1,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(outer.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_0', None, 'test_nested_batch_removes_entire_nested_when_parent_index_removed', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_removal_performance_and_correctness():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    size = 1000
    content = list(range(size))
    indices = [(i,) for i in range(size)]
    batch = Batch(content=content[:], indices=indices[:])
    removal_set: Set[Tuple[int, ...]] = {(i,) for i in range(0, size, 3)}
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=removal_set)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_0', None, 'test_large_scale_removal_performance_and_correctness', 'remove_by_indices', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected_remaining = [i for i in range(size) if i % 3 != 0]