from __future__ import annotations

import gc
import inspect
import os
import threading
import time
from contextlib import contextmanager
from pathlib import Path
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


MODEL_LOCK_ACQUIRE_TIMEOUT = 0.1


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"])
    "Test that _dispose_model_lock removes an existing lock for a model_id."
    manager = ModelManager(ModelRegistry())
    manager._models_state_locks["model_a"] = threading.Lock()
    _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
        "model_a"
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        manager._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
        None,
        "test_dispose_existing_model_lock_removes_lock",
        "_dispose_model_lock",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_nonexistent_model_lock_does_nothing():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test that disposing a lock for a non-existent model_id does nothing."
    manager = ModelManager(ModelRegistry())
    manager._models_state_locks["model_b"] = threading.Lock()
    before = dict(manager._models_state_locks)
    _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
        "model_x"
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        manager._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
        None,
        "test_dispose_nonexistent_model_lock_does_nothing",
        "_dispose_model_lock",
        "4",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_returns_none_when_lock_missing():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test that _dispose_model_lock returns None if the lock is missing."
    manager = ModelManager(ModelRegistry())
    _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
        "model_c"
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        manager._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
        None,
        "test_dispose_model_lock_returns_none_when_lock_missing",
        "_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_locks_dict():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test disposing when _models_state_locks is empty."
    manager = ModelManager(ModelRegistry())
    _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
        "model_z"
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        manager._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
        None,
        "test_dispose_model_lock_with_empty_locks_dict",
        "_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_non_string_model_id():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test disposing with a non-string model_id (should work for any hashable key)."
    manager = ModelManager(ModelRegistry())
    key = 12345
    manager._models_state_locks[key] = threading.Lock()
    _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(key)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        manager._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
        None,
        "test_dispose_model_lock_with_non_string_model_id",
        "_dispose_model_lock",
        "4",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_with_special_characters():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test disposing with a model_id containing special characters."
    manager = ModelManager(ModelRegistry())
    special_id = "model_!@#$%^&*()"
    manager._models_state_locks[special_id] = threading.Lock()
    _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
        special_id
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        manager._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
        None,
        "test_dispose_model_lock_with_special_characters",
        "_dispose_model_lock",
        "4",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_lock_acquisition_failure():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test that if the state lock cannot be acquired, an exception is raised."
    manager = ModelManager(ModelRegistry())
    manager._models_state_locks["model_fail"] = threading.Lock()

    def lock_and_sleep(lock):
        lock.acquire()
        time.sleep(MODEL_LOCK_ACQUIRE_TIMEOUT + 0.2)
        lock.release()

    t = threading.Thread(target=lock_and_sleep, args=(manager._state_lock,))
    t.start()
    time.sleep(0.02)
    with pytest.raises(ModelManagerLockAcquisitionError):
        _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
            "model_fail"
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            manager._dispose_model_lock,
            "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
            None,
            "test_dispose_model_lock_lock_acquisition_failure",
            "_dispose_model_lock",
            "7_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
    t.join()


def test_dispose_model_lock_with_multiple_calls():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test that repeated calls for the same model_id are idempotent."
    manager = ModelManager(ModelRegistry())
    manager._models_state_locks["model_dupe"] = threading.Lock()
    _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
        "model_dupe"
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        manager._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
        None,
        "test_dispose_model_lock_with_multiple_calls",
        "_dispose_model_lock",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
        "model_dupe"
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        manager._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
        None,
        "test_dispose_model_lock_with_multiple_calls",
        "_dispose_model_lock",
        "4",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_with_none_model_id():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test disposing with None as model_id."
    manager = ModelManager(ModelRegistry())
    manager._models_state_locks[None] = threading.Lock()
    _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        manager._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
        None,
        "test_dispose_model_lock_with_none_model_id",
        "_dispose_model_lock",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


def test_dispose_model_lock_with_empty_string_model_id():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test disposing with empty string as model_id."
    manager = ModelManager(ModelRegistry())
    manager._models_state_locks[""] = threading.Lock()
    _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind("")
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        manager._dispose_model_lock,
        "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
        None,
        "test_dispose_model_lock_with_empty_string_model_id",
        "_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"])
    "Test disposing locks in a large manager with many model_ids."
    manager = ModelManager(ModelRegistry())
    num_models = 500
    for i in range(num_models):
        manager._models_state_locks[f"model_{i}"] = threading.Lock()
    for i in range(0, num_models, 2):
        _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
            f"model_{i}"
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            manager._dispose_model_lock,
            "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
            None,
            "test_dispose_model_lock_large_scale_removal",
            "_dispose_model_lock",
            "4_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
    for i in range(num_models):
        if i % 2 == 0:
            pass
        else:
            pass


def test_dispose_model_lock_large_scale_idempotency():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test that disposing all locks twice does not raise and is idempotent."
    manager = ModelManager(ModelRegistry())
    num_models = 300
    for i in range(num_models):
        manager._models_state_locks[f"model_{i}"] = threading.Lock()
    for i in range(num_models):
        _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
            f"model_{i}"
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            manager._dispose_model_lock,
            "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
            None,
            "test_dispose_model_lock_large_scale_idempotency",
            "_dispose_model_lock",
            "4_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
    for i in range(num_models):
        _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
            f"model_{i}"
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            manager._dispose_model_lock,
            "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
            None,
            "test_dispose_model_lock_large_scale_idempotency",
            "_dispose_model_lock",
            "5_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_dispose_model_lock_large_scale_performance():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test performance: disposing locks in a large manager should not be slow."
    manager = ModelManager(ModelRegistry())
    num_models = 800
    for i in range(num_models):
        manager._models_state_locks[f"model_{i}"] = threading.Lock()
    start = time.time()
    for i in range(num_models):
        _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
            f"model_{i}"
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            manager._dispose_model_lock,
            "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
            None,
            "test_dispose_model_lock_large_scale_performance",
            "_dispose_model_lock",
            "5_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
    duration = time.time() - start


def test_dispose_model_lock_large_scale_nonexistent_ids():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test disposing locks for many non-existent model_ids in a large manager."
    manager = ModelManager(ModelRegistry())
    num_models = 100
    for i in range(num_models):
        manager._models_state_locks[f"model_{i}"] = threading.Lock()
    for i in range(num_models, num_models + 100):
        _call__bound__arguments = inspect.signature(manager._dispose_model_lock).bind(
            f"model_{i}"
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            manager._dispose_model_lock,
            "tests.inference.unit_tests.test__dispose_model_lock__unit_test_1",
            None,
            "test_dispose_model_lock_large_scale_nonexistent_ids",
            "_dispose_model_lock",
            "4_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
    for i in range(num_models):
        pass
