from __future__ import annotations

import gc
import os
import time
from typing import Any, Dict, List, Optional
from unittest.mock import MagicMock, Mock, patch

import numpy as np
import pytest
from inference.core.entities.requests.inference import LMMInferenceRequest
from inference.core.entities.responses.inference import InferenceResponse
from inference.core.managers.base import ModelManager
from inference.core.models.base import Model
from inference.core.registries.base import ModelRegistry
from inference.core.workflows.core_steps.common.entities import \
    StepExecutionMode
from inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v1 import \
    Qwen35VLBlockV1
from inference.core.workflows.execution_engine.entities.base import (
    Batch, WorkflowImageData)


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}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

@pytest.fixture
def mock_model_manager():
    """Create a mock ModelManager for testing."""
    manager = Mock(spec=ModelManager)
    manager.add_model = Mock()
    manager.infer_from_request_sync = Mock()
    return manager

@pytest.fixture
def workflow_image_data():
    """Create a WorkflowImageData instance with base64 image."""
    parent_metadata = Mock()
    image_data = WorkflowImageData(parent_metadata=parent_metadata, base64_image='iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==')
    return image_data

@pytest.fixture
def batch_of_images(workflow_image_data):
    """Create a Batch of WorkflowImageData objects."""
    return Batch(content=[workflow_image_data], indices=[(0,)])

@pytest.fixture
def qwen_block_local(mock_model_manager):
    """Create a Qwen35VLBlockV1 instance in LOCAL mode."""
    return Qwen35VLBlockV1(model_manager=mock_model_manager, api_key='test_api_key', step_execution_mode=StepExecutionMode.LOCAL)

@pytest.fixture
def qwen_block_remote(mock_model_manager):
    """Create a Qwen35VLBlockV1 instance in REMOTE mode."""
    return Qwen35VLBlockV1(model_manager=mock_model_manager, api_key='test_api_key', step_execution_mode=StepExecutionMode.REMOTE)

def test_qwen_block_initialization_local(mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test basic initialization of Qwen35VLBlockV1 in LOCAL mode.'
    block = Qwen35VLBlockV1(model_manager=mock_model_manager, api_key='test_key', step_execution_mode=StepExecutionMode.LOCAL)

def test_qwen_block_initialization_remote(mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test basic initialization of Qwen35VLBlockV1 in REMOTE mode.'
    block = Qwen35VLBlockV1(model_manager=mock_model_manager, api_key='test_key', step_execution_mode=StepExecutionMode.REMOTE)

def test_run_local_execution_mode(qwen_block_local, batch_of_images, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that run() calls run_locally when execution mode is LOCAL.'
    qwen_block_local.run_locally = Mock(return_value=[{'parsed_output': 'test'}])
    codeflash_return_value = codeflash_wrap(qwen_block_local.run, 'tests.inference.unit_tests.test_run__unit_test_1', None, 'test_run_local_execution_mode', 'run', '2', codeflash_loop_index, images=batch_of_images, model_version='qwen3.5-vl', prompt='Test prompt', system_prompt='Test system')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    qwen_block_local.run_locally.assert_called_once()

def test_run_remote_execution_mode(qwen_block_remote, batch_of_images):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that run() calls run_remotely when execution mode is REMOTE.'
    qwen_block_remote.run_remotely = Mock(return_value=[{'parsed_output': 'test'}])
    codeflash_return_value = codeflash_wrap(qwen_block_remote.run, 'tests.inference.unit_tests.test_run__unit_test_1', None, 'test_run_remote_execution_mode', 'run', '2', codeflash_loop_index, images=batch_of_images, model_version='qwen3.5-vl', prompt='Test prompt', system_prompt='Test system')
    codeflash_output = codeflash_return_value
    result = codeflash_output
    qwen_block_remote.run_remotely.assert_called_once()

def test_run_invalid_execution_mode(mock_model_manager, batch_of_images):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that run() raises ValueError for invalid execution mode.'
    block = Qwen35VLBlockV1(model_manager=mock_model_manager, api_key='test_key', step_execution_mode=StepExecutionMode.LOCAL)
    block._step_execution_mode = 'invalid_mode'
    with pytest.raises(ValueError, match='Unknown step execution mode'):
        codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_1', None, 'test_run_invalid_execution_mode', 'run', '3_0', codeflash_loop_index, images=batch_of_images, model_version='qwen3.5-vl', prompt='Test', system_prompt='Test')

def test_run_locally_default_prompts(qwen_block_local, batch_of_images, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test run_locally with default prompts when None is provided.'
    mock_response = Mock(spec=InferenceResponse)
    mock_response.response = 'Generated response'
    mock_model_manager.infer_from_request_sync.return_value = mock_response
    result = qwen_block_local.run_locally(images=batch_of_images, model_version='qwen3.5-vl', prompt=None, system_prompt=None)
    mock_model_manager.add_model.assert_called_once_with(model_id='qwen3.5-vl', api_key='test_api_key')

def test_run_locally_custom_prompts(qwen_block_local, batch_of_images, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test run_locally with custom prompts.'
    mock_response = Mock(spec=InferenceResponse)
    mock_response.response = 'Custom response'
    mock_model_manager.infer_from_request_sync.return_value = mock_response
    custom_prompt = 'What is in this image?'
    custom_system = 'You are a helpful assistant.'
    result = qwen_block_local.run_locally(images=batch_of_images, model_version='qwen3.5-vl', prompt=custom_prompt, system_prompt=custom_system)

def test_run_locally_empty_batch(qwen_block_local, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test run_locally with empty batch of images.'
    empty_batch = Batch(content=[], indices=[])
    result = qwen_block_local.run_locally(images=empty_batch, model_version='qwen3.5-vl', prompt='Test', system_prompt='Test')
    mock_model_manager.add_model.assert_called_once()

def test_run_locally_none_api_key(mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test run_locally with None api_key.'
    block = Qwen35VLBlockV1(model_manager=mock_model_manager, api_key=None, step_execution_mode=StepExecutionMode.LOCAL)
    parent_metadata = Mock()
    image = WorkflowImageData(parent_metadata=parent_metadata, base64_image='iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==')
    batch = Batch(content=[image], indices=[(0,)])
    mock_response = Mock(spec=InferenceResponse)
    mock_response.response = 'Test'
    mock_model_manager.infer_from_request_sync.return_value = mock_response
    result = block.run_locally(images=batch, model_version='qwen3.5-vl', prompt='Test', system_prompt='Test')

def test_run_locally_empty_string_prompts(qwen_block_local, batch_of_images, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test run_locally with empty string prompts (should use defaults).'
    mock_response = Mock(spec=InferenceResponse)
    mock_response.response = 'Response'
    mock_model_manager.infer_from_request_sync.return_value = mock_response
    result = qwen_block_local.run_locally(images=batch_of_images, model_version='qwen3.5-vl', prompt='', system_prompt='')
    call_args = mock_model_manager.infer_from_request_sync.call_args
    request = call_args.kwargs['request']

def test_run_locally_special_characters_in_prompt(qwen_block_local, batch_of_images, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test run_locally with special characters in prompts.'
    mock_response = Mock(spec=InferenceResponse)
    mock_response.response = 'Response'
    mock_model_manager.infer_from_request_sync.return_value = mock_response
    special_prompt = 'What is <this> & \'that\'? "Quote"'
    special_system = 'You are <a> model!'
    result = qwen_block_local.run_locally(images=batch_of_images, model_version='qwen3.5-vl', prompt=special_prompt, system_prompt=special_system)
    call_args = mock_model_manager.infer_from_request_sync.call_args
    request = call_args.kwargs['request']

def test_run_locally_multiple_images(qwen_block_local, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test run_locally with multiple images in batch.'
    parent_metadata = Mock()
    image1 = WorkflowImageData(parent_metadata=parent_metadata, base64_image='iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==')
    image2 = WorkflowImageData(parent_metadata=parent_metadata, base64_image='iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==')
    batch = Batch(content=[image1, image2], indices=[(0,), (1,)])
    mock_response1 = Mock(spec=InferenceResponse)
    mock_response1.response = 'Response 1'
    mock_response2 = Mock(spec=InferenceResponse)
    mock_response2.response = 'Response 2'
    mock_model_manager.infer_from_request_sync.side_effect = [mock_response1, mock_response2]
    result = qwen_block_local.run_locally(images=batch, model_version='qwen3.5-vl', prompt='Test', system_prompt='Test')

def test_run_locally_lmm_request_structure(qwen_block_local, batch_of_images, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that run_locally creates proper LMMInferenceRequest objects.'
    mock_response = Mock(spec=InferenceResponse)
    mock_response.response = 'Test response'
    mock_model_manager.infer_from_request_sync.return_value = mock_response
    result = qwen_block_local.run_locally(images=batch_of_images, model_version='qwen3.5-vl-7b', prompt='Describe this', system_prompt='Be concise')
    call_args = mock_model_manager.infer_from_request_sync.call_args
    request = call_args.kwargs['request']

def test_run_locally_prompt_concatenation(qwen_block_local, batch_of_images, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that prompts are correctly concatenated with <system_prompt> marker.'
    mock_response = Mock(spec=InferenceResponse)
    mock_response.response = 'Test'
    mock_model_manager.infer_from_request_sync.return_value = mock_response
    prompt = 'What do you see?'
    system = 'Answer concisely.'
    qwen_block_local.run_locally(images=batch_of_images, model_version='qwen3.5-vl', prompt=prompt, system_prompt=system)
    call_args = mock_model_manager.infer_from_request_sync.call_args
    request = call_args.kwargs['request']
    expected_prompt = f'{prompt}<system_prompt>{system}'

def test_run_locally_to_inference_format_called(qwen_block_local, batch_of_images, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that to_inference_format is called on images.'
    mock_response = Mock(spec=InferenceResponse)
    mock_response.response = 'Test'
    mock_model_manager.infer_from_request_sync.return_value = mock_response
    original_image = batch_of_images._content[0]
    original_image.to_inference_format = Mock(return_value={'type': 'base64', 'value': 'test_base64'})
    qwen_block_local.run_locally(images=batch_of_images, model_version='qwen3.5-vl', prompt='Test', system_prompt='Test')
    original_image.to_inference_format.assert_called_once_with(numpy_preferred=False)

def test_run_locally_large_batch_50_images(qwen_block_local, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test run_locally with a large batch of 50 images.'
    parent_metadata = Mock()
    images = []
    for i in range(50):
        image = WorkflowImageData(parent_metadata=parent_metadata, base64_image='iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==')
        images.append(image)
    batch = Batch(content=images, indices=[(i,) for i in range(50)])
    mock_responses = []
    for i in range(50):
        mock_response = Mock(spec=InferenceResponse)
        mock_response.response = f'Response {i}'
        mock_responses.append(mock_response)
    mock_model_manager.infer_from_request_sync.side_effect = mock_responses
    result = qwen_block_local.run_locally(images=batch, model_version='qwen3.5-vl', prompt='Test', system_prompt='Test')
    for i in range(50):
        pass

def test_run_locally_large_prompts(qwen_block_local, batch_of_images, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test run_locally with very large prompts (1000+ characters).'
    mock_response = Mock(spec=InferenceResponse)
    mock_response.response = 'Test response'
    mock_model_manager.infer_from_request_sync.return_value = mock_response
    large_prompt = 'What do you see? ' * 100
    large_system = 'Be detailed. ' * 100
    result = qwen_block_local.run_locally(images=batch_of_images, model_version='qwen3.5-vl', prompt=large_prompt, system_prompt=large_system)
    call_args = mock_model_manager.infer_from_request_sync.call_args
    request = call_args.kwargs['request']

def test_run_locally_batch_100_images_stress(qwen_block_local, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Stress test with 100 images to ensure scalability.'
    parent_metadata = Mock()
    images = []
    for i in range(100):
        image = WorkflowImageData(parent_metadata=parent_metadata, base64_image='iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==')
        images.append(image)
    batch = Batch(content=images, indices=[(i,) for i in range(100)])

    def create_response(i):
        mock_response = Mock(spec=InferenceResponse)
        mock_response.response = f'Response {i}'
        return mock_response
    mock_model_manager.infer_from_request_sync.side_effect = [create_response(i) for i in range(100)]
    result = qwen_block_local.run_locally(images=batch, model_version='qwen3.5-vl', prompt='Analyze', system_prompt='Concise')
    for i in range(100):
        pass

def test_run_locally_repeated_calls_with_state(qwen_block_local, batch_of_images, mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that repeated calls to run_locally work correctly (state persistence).'
    mock_response = Mock(spec=InferenceResponse)
    mock_response.response = 'Response'
    mock_model_manager.infer_from_request_sync.return_value = mock_response
    for i in range(10):
        result = qwen_block_local.run_locally(images=batch_of_images, model_version='qwen3.5-vl', prompt=f'Test {i}', system_prompt=f'System {i}')

def test_run_remotely_result_format(qwen_block_remote, batch_of_images):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that run_remotely returns results in the correct format.'
    with patch('inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v1.InferenceHTTPClient') as mock_client_class:
        mock_client = Mock()
        mock_client_class.return_value = mock_client
        mock_client.infer_lmm.return_value = {'response': 'Test response'}
        result = qwen_block_remote.run_remotely(images=batch_of_images, model_version='qwen3.5-vl', prompt='Test', system_prompt='Test')

def test_run_remotely_multiple_images_large_batch(qwen_block_remote):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test run_remotely with a large batch of images.'
    parent_metadata = Mock()
    images = []
    for i in range(20):
        image = WorkflowImageData(parent_metadata=parent_metadata, base64_image='iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==')
        images.append(image)
    batch = Batch(content=images, indices=[(i,) for i in range(20)])
    with patch('inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v1.InferenceHTTPClient') as mock_client_class:
        mock_client = Mock()
        mock_client_class.return_value = mock_client
        mock_client.infer_lmm.return_value = {'response': 'Response'}
        result = qwen_block_remote.run_remotely(images=batch, model_version='qwen3.5-vl', prompt='Test', system_prompt='Test')

def test_run_integration_local_remote_consistency(mock_model_manager):
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that local and remote execution produce consistent result structure.'
    parent_metadata = Mock()
    image = WorkflowImageData(parent_metadata=parent_metadata, base64_image='iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==')
    batch = Batch(content=[image], indices=[(0,)])
    block_local = Qwen35VLBlockV1(model_manager=mock_model_manager, api_key='test_key', step_execution_mode=StepExecutionMode.LOCAL)
    mock_response = Mock(spec=InferenceResponse)
    mock_response.response = 'Test response'
    mock_model_manager.infer_from_request_sync.return_value = mock_response
    result_local = block_local.run_locally(images=batch, model_version='qwen3.5-vl', prompt='Test', system_prompt='Test')
    block_remote = Qwen35VLBlockV1(model_manager=mock_model_manager, api_key='test_key', step_execution_mode=StepExecutionMode.REMOTE)
    with patch('inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v1.InferenceHTTPClient') as mock_client_class:
        mock_client = Mock()
        mock_client_class.return_value = mock_client
        mock_client.infer_lmm.return_value = {'response': 'Test response'}
        result_remote = block_remote.run_remotely(images=batch, model_version='qwen3.5-vl', prompt='Test', system_prompt='Test')