from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path

import dill as pickle
import inference.core.cache as cache_module
import inference.core.env as env_module
import pytest
from inference.core.cache import cache
from inference.core.cache.memory import MemoryCache
from inference.core.devices.utils import GLOBAL_INFERENCE_SERVER_ID
from inference.core.entities.requests.inference import InferenceRequest
from inference.core.entities.responses.inference import InferenceResponse
from inference.core.exceptions import InferenceModelNotFound
from inference.core.managers.base import ModelManager
from inference.core.models.base import Model as BaseModel
from inference.core.registries.base import ModelRegistry


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

class TestModel(BaseModel):
    """A minimal concrete model that returns a simple InferenceResponse."""

    def infer(self, **kwargs):
        resp = InferenceResponse()
        return [resp]

class TestModelRaises(BaseModel):
    """A minimal concrete model that raises during inference to test error-handling."""

    def infer_from_request(self, request):
        raise RuntimeError('intentional inference failure for testing')

def test_infer_from_request_sync_returns_response_and_caches_request():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    registry = ModelRegistry({})
    model_instance = TestModel()
    manager = ModelManager(model_registry=registry, models={'m1': model_instance})
    request = InferenceRequest(api_key='test-api-key', image='raw-image-data', model_id='m1')
    _call__bound__arguments = inspect.signature(manager.infer_from_request_sync).bind('m1', request)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager.infer_from_request_sync, 'tests.inference.unit_tests.test_infer_from_request_sync__unit_test_0', None, 'test_infer_from_request_sync_returns_response_and_caches_request', 'infer_from_request_sync', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    cached_keys = set(cache_module.cache.cache.keys())
    expected_inference_prefix = f'inference:{GLOBAL_INFERENCE_SERVER_ID}:m1'
    has_inference_key = any((k.startswith(expected_inference_prefix) for k in cached_keys))

def test_infer_from_request_sync_raises_when_model_not_loaded():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    registry = ModelRegistry({})
    manager = ModelManager(model_registry=registry, models={})
    request = InferenceRequest(api_key='k', image='img', model_id='missing')
    with pytest.raises(InferenceModelNotFound):
        _call__bound__arguments = inspect.signature(manager.infer_from_request_sync).bind('missing', request)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(manager.infer_from_request_sync, 'tests.inference.unit_tests.test_infer_from_request_sync__unit_test_0', None, 'test_infer_from_request_sync_raises_when_model_not_loaded', 'infer_from_request_sync', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_exception_during_inference_is_propagated_and_error_cached():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    registry = ModelRegistry({})
    failing_model = TestModelRaises()
    manager = ModelManager(model_registry=registry, models={'bad': failing_model})
    request = InferenceRequest(api_key='api-e', image='img', model_id='bad')
    with pytest.raises(RuntimeError) as excinfo:
        _call__bound__arguments = inspect.signature(manager.infer_from_request_sync).bind('bad', request)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(manager.infer_from_request_sync, 'tests.inference.unit_tests.test_infer_from_request_sync__unit_test_0', None, 'test_exception_during_inference_is_propagated_and_error_cached', 'infer_from_request_sync', '4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    cached_keys = set(cache_module.cache.cache.keys())
    expected_error_prefix = f'error:{GLOBAL_INFERENCE_SERVER_ID}:bad'
    has_error_key = any((k.startswith(expected_error_prefix) for k in cached_keys))

def test_disable_model_monitoring_prevents_caching_even_when_inference_succeeds():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    registry = ModelRegistry({})
    model_instance = TestModel()
    manager = ModelManager(model_registry=registry, models={'m2': model_instance})
    request = InferenceRequest(api_key='akey', image='img', model_id='m2', disable_model_monitoring=True)
    _call__bound__arguments = inspect.signature(manager.infer_from_request_sync).bind('m2', request)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager.infer_from_request_sync, 'tests.inference.unit_tests.test_infer_from_request_sync__unit_test_0', None, 'test_disable_model_monitoring_prevents_caching_even_when_inference_succeeds', 'infer_from_request_sync', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    cached_keys = set(cache_module.cache.cache.keys())
    inference_prefix = f'inference:{GLOBAL_INFERENCE_SERVER_ID}:m2'

def test_numpy_image_value_is_stringified_before_caching():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

    class ImageLike:

        def __init__(self, t, v):
            self.type = t
            self.value = v
    registry = ModelRegistry({})
    model_instance = TestModel()
    manager = ModelManager(model_registry=registry, models={'m3': model_instance})
    numpy_image = ImageLike('numpy', 12345)
    request = InferenceRequest(api_key='k3', image=numpy_image, model_id='m3')
    _call__bound__arguments = inspect.signature(manager.infer_from_request_sync).bind('m3', request)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager.infer_from_request_sync, 'tests.inference.unit_tests.test_infer_from_request_sync__unit_test_0', None, 'test_numpy_image_value_is_stringified_before_caching', 'infer_from_request_sync', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    _ = codeflash_output
    cached = cache_module.cache.cache
    inference_key = next((k for k in cached.keys() if k.startswith(f'inference:{GLOBAL_INFERENCE_SERVER_ID}:m3')), None)
    inner = cached[inference_key]
    any_value = next(iter(inner.values()))
    serialized = str(any_value)

def test_large_scale_many_inferences_cached_and_retrievable():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "\n    Run many inferences sequentially to exercise caching logic at scale (up to 1000 items).\n    We perform 1000 fast synthetic inferences and verify that the store contains\n    the expected number of model entries in the 'inference' sorted-set.\n    "
    registry = ModelRegistry({})
    model_instance = TestModel()
    manager = ModelManager(model_registry=registry, models={'bulk': model_instance})
    n = 1000
    for i in range(n):
        req = InferenceRequest(api_key=f'bulk-key-{i}', image=f'img-{i}', model_id='bulk')
        _call__bound__arguments = inspect.signature(manager.infer_from_request_sync).bind('bulk', req)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(manager.infer_from_request_sync, 'tests.inference.unit_tests.test_infer_from_request_sync__unit_test_0', None, 'test_large_scale_many_inferences_cached_and_retrievable', 'infer_from_request_sync', '5_1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    cached = cache_module.cache.cache
    inference_key = next((k for k in cached.keys() if k.startswith(f'inference:{GLOBAL_INFERENCE_SERVER_ID}:bulk')), None)
    inner = cached[inference_key]