from __future__ import annotations

import gc
import os
import time

import pytest
import torch
from inference.models.qwen3_5vl.qwen3_5vl_inference_models import \
    InferenceModelsQwen35VLAdapter
from inference_models.models.qwen3_5.qwen3_5_hf import Qwen35HF


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
    _codeflash_should_sync_cuda = torch.cuda.is_available() and torch.cuda.is_initialized()
    _codeflash_should_sync_mps = not _codeflash_should_sync_cuda and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() and hasattr(torch.mps, 'synchronize')
    gc.disable()
    try:
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        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}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

def _make_adapter_with_generate(generate_fn):
    adapter = InferenceModelsQwen35VLAdapter.__new__(InferenceModelsQwen35VLAdapter)
    adapter.metrics = {'num_inferences': 0, 'avg_inference_time': 0.0}
    adapter.api_key = 'test-api-key'
    adapter.task_type = 'lmm'
    model_instance = Qwen35HF.__new__(Qwen35HF)
    model_instance.generate = generate_fn
    adapter._model = model_instance
    return adapter
CANONICAL_INPUT = torch.tensor([1.0, -2.0, 3.5, 0.0], dtype=torch.float32)

def test_predict_forwards_inputs_and_kwargs_basic():
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Basic functionality:\n    - The adapter.predict should call underlying model.generate with the given inputs\n      and keyword arguments (after mapping).\n    - The return value should be whatever the underlying generate returns.\n    '

    def generate_fn(inputs, **kwargs):
        scale = kwargs.get('scale', 1.0)
        return inputs * float(scale)
    adapter = _make_adapter_with_generate(generate_fn)
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_predict_forwards_inputs_and_kwargs_basic', 'predict', '3', codeflash_loop_index, CANONICAL_INPUT, scale=2.5)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = CANONICAL_INPUT * 2.5

def test_predict_uses_map_inference_kwargs_override():
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "\n    Verify that map_inference_kwargs is used to transform kwargs before calling generate.\n    We'll override map_inference_kwargs on the adapter instance to add/modify kwargs.\n    "

    def generate_fn(inputs, **kwargs):
        scale = kwargs.get('scale', 1.0)
        return inputs * float(scale)
    adapter = _make_adapter_with_generate(generate_fn)

    def mapper(kwargs: dict) -> dict:
        new = dict(kwargs)
        new['scale'] = 3.0
        return new
    adapter.map_inference_kwargs = mapper
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_predict_uses_map_inference_kwargs_override', 'predict', '5', codeflash_loop_index, CANONICAL_INPUT)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = CANONICAL_INPUT * 3.0

def test_edge_cases_various_input_shapes_and_values():
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Edge cases are grouped here (per JIT rules) to avoid recompilations:\n    - empty tensor (0,)\n    - tensor containing NaN and Inf\n    - canonical non-empty tensor (re-check)\n    '

    def generate_fn(inputs, **kwargs):
        return inputs.clone()
    adapter = _make_adapter_with_generate(generate_fn)
    empty = torch.empty((0,), dtype=torch.float32)
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_edge_cases_various_input_shapes_and_values', 'predict', '4', codeflash_loop_index, empty)
    codeflash_output = codeflash_return_value
    out_empty = codeflash_output
    nan_inf = torch.tensor([float('nan'), float('inf'), -float('inf'), 5.0], dtype=torch.float32)
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_edge_cases_various_input_shapes_and_values', 'predict', '7', codeflash_loop_index, nan_inf)
    codeflash_output = codeflash_return_value
    out_nan_inf = codeflash_output
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_edge_cases_various_input_shapes_and_values', 'predict', '9', codeflash_loop_index, CANONICAL_INPUT)
    codeflash_output = codeflash_return_value
    out_canon = codeflash_output

def test_predict_many_iterations_consistent_and_counted():
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Large-scale style test (many repeated calls, up to 1000 iterations).\n    - Uses the canonical input shape/dtype for every call (avoids recompilation).\n    - Ensures that repeated calls produce consistent results.\n    - Verifies number of generate invocations equals number of predict calls.\n    '
    call_count = {'count': 0}

    def generate_fn(inputs, **kwargs):
        call_count['count'] += 1
        return inputs + 1.0
    adapter = _make_adapter_with_generate(generate_fn)
    iterations = 1000
    for _ in range(iterations):
        codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_predict_many_iterations_consistent_and_counted', 'predict', '5_0', codeflash_loop_index, CANONICAL_INPUT)
        codeflash_output = codeflash_return_value
        out = codeflash_output
        expected = CANONICAL_INPUT + 1.0