from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path
from typing import Any, Callable, Optional

import dill as pickle
import pytest
from inference.core.workflows.core_steps.analytics.detection_event_log.v1 import (
    MAX_VIDEOS, DetectionEventLogBlockV1)


def codeflash_wrap(
    codeflash_wrapped: Callable[..., Any],
    codeflash_test_module_name: str,
    codeflash_test_class_name: str | None,
    codeflash_test_name: str,
    codeflash_function_name: str,
    codeflash_line_id: str,
    codeflash_loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> Any:
    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_no_eviction_when_under_max_videos():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    for i in range(MAX_VIDEOS):
        vid = f'video_{i}'
        block._event_logs[vid] = {'dummy': 'data'}
        block._last_flush_frame[vid] = 0
        block._frame_count[vid] = 0
        block._last_access[vid] = i
    _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_no_eviction_when_under_max_videos', '_evict_oldest_video', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_eviction_when_exceeding_max_videos():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    for i in range(MAX_VIDEOS + 1):
        vid = f'video_{i}'
        block._event_logs[vid] = {'dummy': 'data'}
        block._last_flush_frame[vid] = i
        block._frame_count[vid] = i
        block._last_access[vid] = i
    _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_eviction_when_exceeding_max_videos', '_evict_oldest_video', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    for i in range(1, MAX_VIDEOS + 1):
        vid = f'video_{i}'

def test_eviction_removes_all_video_data():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    block._event_logs['video_a'] = {'dummy': 'data'}
    block._event_logs['video_b'] = {'dummy': 'data'}
    block._last_flush_frame['video_a'] = 10
    block._last_flush_frame['video_b'] = 20
    block._frame_count['video_a'] = 5
    block._frame_count['video_b'] = 15
    block._last_access['video_a'] = 1
    block._last_access['video_b'] = 2
    block._event_logs['video_c'] = {'dummy': 'data'}
    block._last_flush_frame['video_c'] = 30
    block._frame_count['video_c'] = 25
    block._last_access['video_c'] = 3
    _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_eviction_removes_all_video_data', '_evict_oldest_video', '13', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    for vid in ['video_b', 'video_c']:
        pass

def test_eviction_with_empty_dicts():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_eviction_with_empty_dicts', '_evict_oldest_video', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_eviction_with_missing_keys():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    block._event_logs['video_x'] = {'dummy': 'data'}
    block._last_access['video_x'] = 42
    block._event_logs['video_y'] = {'dummy': 'data'}
    block._last_access['video_y'] = 43
    block._event_logs['video_z'] = {'dummy': 'data'}
    block._last_access['video_z'] = 44
    _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_eviction_with_missing_keys', '_evict_oldest_video', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_eviction_with_duplicate_last_access_times():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    block._event_logs['video_1'] = {'dummy': 'data'}
    block._event_logs['video_2'] = {'dummy': 'data'}
    block._event_logs['video_3'] = {'dummy': 'data'}
    block._last_access['video_1'] = 100
    block._last_access['video_2'] = 100
    block._last_access['video_3'] = 101
    block._last_flush_frame['video_1'] = 1
    block._last_flush_frame['video_2'] = 2
    block._last_flush_frame['video_3'] = 3
    block._frame_count['video_1'] = 10
    block._frame_count['video_2'] = 20
    block._frame_count['video_3'] = 30
    _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_eviction_with_duplicate_last_access_times', '_evict_oldest_video', '13', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    evicted = {'video_1', 'video_2'} - set(block._event_logs.keys())
    remaining = set(block._event_logs.keys())
    evicted_vid = evicted.pop()

def test_eviction_with_non_integer_last_access():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    block._event_logs['video_1'] = {'dummy': 'data'}
    block._event_logs['video_2'] = {'dummy': 'data'}
    block._last_access['video_1'] = 100.5
    block._last_access['video_2'] = 99.9
    _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_eviction_with_non_integer_last_access', '_evict_oldest_video', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_eviction_with_missing_last_access():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    block._event_logs['video_1'] = {'dummy': 'data'}
    block._event_logs['video_2'] = {'dummy': 'data'}
    block._last_access['video_1'] = 1
    block._event_logs['video_3'] = {'dummy': 'data'}
    block._last_access['video_3'] = 2
    with pytest.raises(ValueError):
        _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_eviction_with_missing_last_access', '_evict_oldest_video', '6_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_scale_eviction_performance():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    N = 1000
    for i in range(N):
        vid = f'video_{i}'
        block._event_logs[vid] = {'dummy': 'data'}
        block._last_flush_frame[vid] = i
        block._frame_count[vid] = i
        block._last_access[vid] = i
    block._event_logs['video_extra'] = {'dummy': 'data'}
    block._last_flush_frame['video_extra'] = N
    block._frame_count['video_extra'] = N
    block._last_access['video_extra'] = N
    _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_large_scale_eviction_performance', '_evict_oldest_video', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    for i in range(1, N):
        vid = f'video_{i}'

def test_large_scale_eviction_multiple_times():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    N = MAX_VIDEOS + 10
    for i in range(N):
        vid = f'video_{i}'
        block._event_logs[vid] = {'dummy': 'data'}
        block._last_flush_frame[vid] = i
        block._frame_count[vid] = i
        block._last_access[vid] = i
    for _ in range(10):
        _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_large_scale_eviction_multiple_times', '_evict_oldest_video', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    for i in range(10):
        vid = f'video_{i}'
    for i in range(10, N):
        vid = f'video_{i}'

def test_large_scale_eviction_with_sparse_last_access():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    N = MAX_VIDEOS + 5
    for i in range(N):
        vid = f'video_{i}'
        block._event_logs[vid] = {'dummy': 'data'}
        block._last_flush_frame[vid] = i
        block._frame_count[vid] = i
        block._last_access[vid] = i * 10
    _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_large_scale_eviction_with_sparse_last_access', '_evict_oldest_video', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    for i in range(1, N):
        vid = f'video_{i}'

def test_massive_scale_eviction_5000_videos():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = DetectionEventLogBlockV1()
    N = 5000
    for i in range(N):
        vid = f'video_{i}'
        block._event_logs[vid] = {'dummy': 'data'}
        block._last_flush_frame[vid] = i
        block._frame_count[vid] = i
        block._last_access[vid] = i
    for evict_round in range(50):
        _call__bound__arguments = inspect.signature(block._evict_oldest_video).bind()
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(block._evict_oldest_video, 'tests.inference.unit_tests.test__evict_oldest_video__unit_test_0', None, 'test_massive_scale_eviction_5000_videos', '_evict_oldest_video', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)