from __future__ import annotations

import gc
import inspect
import os
import threading
import time
from contextlib import contextmanager
from pathlib import Path
from threading import Lock
from typing import Any, Callable, Dict, Generator, Optional

import dill as pickle
import pytest
from inference.core.managers.base import ModelManager


def codeflash_wrap(
    wrapped: Callable[..., Any],
    test_module_name: str,
    test_class_name: str | None,
    test_name: str,
    function_name: str,
    line_id: str,
    loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> Any:
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    test_stdout_tag = f"{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}"
    print(f"!$######{test_stdout_tag}######$!")
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = 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"{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{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(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 ModelManagerLockAcquisitionError(Exception):
    pass


class Model:
    pass


class ModelRegistry:
    pass


def test_dispose_existing_model_lock_removes_lock():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    mm._models_state_locks["model_a"] = Lock()
    _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind("model_a")
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_existing_model_lock_removes_lock",
        "_dispose_model_lock",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_nonexistent_model_lock_does_nothing():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind("model_x")
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_nonexistent_model_lock_does_nothing",
        "_dispose_model_lock",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_dispose_model_lock_with_multiple_locks_removes_only_target():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    mm._models_state_locks["model_a"] = Lock()
    mm._models_state_locks["model_b"] = Lock()
    mm._models_state_locks["model_c"] = Lock()
    _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind("model_b")
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_model_lock_with_multiple_locks_removes_only_target",
        "_dispose_model_lock",
        "4",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_returns_none_when_lock_removed():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    mm._models_state_locks["model_a"] = Lock()
    _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind("model_a")
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_model_lock_returns_none_when_lock_removed",
        "_dispose_model_lock",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_dispose_model_lock_with_empty_dict():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind(
        "any_model"
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_model_lock_with_empty_dict",
        "_dispose_model_lock",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_dispose_model_lock_with_model_id_not_string():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    mm._models_state_locks[123] = Lock()
    _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind(123)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_model_lock_with_model_id_not_string",
        "_dispose_model_lock",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_with_special_characters_in_model_id():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    mm._models_state_locks["!@#$_model"] = Lock()
    _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind(
        "!@#$_model"
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_model_lock_with_special_characters_in_model_id",
        "_dispose_model_lock",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_when_lock_cannot_be_acquired():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    mm._models_state_locks["model_a"] = Lock()
    mm._state_lock.acquire()
    try:
        with pytest.raises(ModelManagerLockAcquisitionError):
            _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind(
                "model_a"
            )
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(
                mm._dispose_model_lock,
                "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
                None,
                "test_dispose_model_lock_when_lock_cannot_be_acquired",
                "_dispose_model_lock",
                "3_0_0",
                codeflash_loop_index,
                **_call__bound__arguments.arguments,
            )
    finally:
        mm._state_lock.release()


def test_dispose_model_lock_with_none_model_id():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    mm._models_state_locks[None] = Lock()
    _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind(None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_model_lock_with_none_model_id",
        "_dispose_model_lock",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_with_large_string_model_id():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    long_id = "model_" + "x" * 500
    mm = ModelManager(ModelRegistry())
    mm._models_state_locks[long_id] = Lock()
    _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind(long_id)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_model_lock_with_large_string_model_id",
        "_dispose_model_lock",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_with_model_id_collision():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    mm._models_state_locks["model"] = Lock()
    mm._models_state_locks["model "] = Lock()
    _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind("model ")
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_model_lock_with_model_id_collision",
        "_dispose_model_lock",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_large_scale_removal():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    for i in range(1000):
        mm._models_state_locks[f"model_{i}"] = Lock()
    for i in range(0, 1000, 100):
        _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind(
            f"model_{i}"
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            mm._dispose_model_lock,
            "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
            None,
            "test_dispose_model_lock_large_scale_removal",
            "_dispose_model_lock",
            "2_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
    for i in range(0, 1000, 100):
        pass
    for i in range(1, 1000):
        if i % 100 != 0:
            pass


def test_dispose_model_lock_performance_under_load():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    for i in range(500):
        mm._models_state_locks[f"model_{i}"] = Lock()
    start = time.time()
    for i in range(500):
        _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind(
            f"model_{i}"
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            mm._dispose_model_lock,
            "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
            None,
            "test_dispose_model_lock_performance_under_load",
            "_dispose_model_lock",
            "3_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
    duration = time.time() - start


def test_dispose_model_lock_thread_safety_under_contention():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    mm._models_state_locks["model_a"] = Lock()

    def hold_lock(lock, duration):
        lock.acquire()
        time.sleep(duration)
        lock.release()

    t = threading.Thread(target=hold_lock, args=(mm._state_lock, 0.2))
    t.start()
    time.sleep(0.05)
    with pytest.raises(ModelManagerLockAcquisitionError):
        _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind(
            "model_a"
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            mm._dispose_model_lock,
            "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
            None,
            "test_dispose_model_lock_thread_safety_under_contention",
            "_dispose_model_lock",
            "6_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
    t.join()
    _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind("model_a")
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_model_lock_thread_safety_under_contention",
        "_dispose_model_lock",
        "8",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_removal_does_not_affect_other_manager_instances():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm1 = ModelManager(ModelRegistry())
    mm2 = ModelManager(ModelRegistry())
    mm1._models_state_locks["model_a"] = Lock()
    mm2._models_state_locks["model_a"] = Lock()
    _call__bound__arguments = inspect.signature(mm1._dispose_model_lock).bind("model_a")
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        mm1._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
        None,
        "test_dispose_model_lock_removal_does_not_affect_other_manager_instances",
        "_dispose_model_lock",
        "4",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_large_scale_with_varied_keys():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mm = ModelManager(ModelRegistry())
    for i in range(500):
        mm._models_state_locks[i] = Lock()
    for i in range(500):
        mm._models_state_locks[f"model_{i}"] = Lock()
    for i in range(500):
        _call__bound__arguments = inspect.signature(mm._dispose_model_lock).bind(i)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            mm._dispose_model_lock,
            "tests.inference.unit_tests.test__dispose_model_lock__unit_test_0",
            None,
            "test_dispose_model_lock_large_scale_with_varied_keys",
            "_dispose_model_lock",
            "3_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
    for i in range(500):
        pass
