from __future__ import annotations

import gc
import importlib
import os
import time
from typing import List, Tuple

import inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v1 as qwen_module
import pytest
from inference.core.env import HOSTED_CORE_MODEL_URL
from inference.core.managers.base import ModelManager
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)
from inference_sdk import InferenceHTTPClient


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

def test_run_dispatches_to_local_when_mode_is_local():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_manager = ModelManager(ModelRegistry({}))
    block = Qwen35VLBlockV1(model_manager=model_manager, api_key='test-key', step_execution_mode=StepExecutionMode.LOCAL)
    img = WorkflowImageData(parent_metadata=object(), base64_image='b64-1')
    batch = Batch.init(content=[img], indices=[(0,)])
    sentinel = [{'parsed_output': 'LOCAL_OK'}]
    block.run_locally = lambda images, model_version, prompt, system_prompt: sentinel
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_dispatches_to_local_when_mode_is_local', 'run', '6', codeflash_loop_index, images=batch, model_version='m:v', prompt=None, system_prompt=None)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_run_dispatches_to_remote_when_mode_is_remote(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_manager = ModelManager(ModelRegistry({}))
    block = Qwen35VLBlockV1(model_manager=model_manager, api_key='rk_test', step_execution_mode=StepExecutionMode.REMOTE)
    img = WorkflowImageData(parent_metadata=object(), base64_image='b64-remote')
    batch = Batch.init(content=[img], indices=[(0,)])
    captured = {}

    def fake_infer_lmm(self, inference_input, model_id, prompt=None, model_id_in_path=False):
        captured['inference_input'] = inference_input
        captured['model_id'] = model_id
        captured['prompt'] = prompt
        captured['model_id_in_path'] = model_id_in_path
        return {'response': f'resp-for-{inference_input}'}
    monkeypatch.setattr(InferenceHTTPClient, 'infer_lmm', fake_infer_lmm, raising=True)
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_dispatches_to_remote_when_mode_is_remote', 'run', '7', codeflash_loop_index, images=batch, model_version='some-model', prompt='Q?', system_prompt='SYS')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_run_raises_value_error_for_unknown_execution_mode():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_manager = ModelManager(ModelRegistry({}))
    block = Qwen35VLBlockV1(model_manager=model_manager, api_key=None, step_execution_mode='weird-mode')
    img = WorkflowImageData(parent_metadata=object(), base64_image='b64-x')
    batch = Batch.init(content=[img], indices=[(0,)])
    with pytest.raises(ValueError) as excinfo:
        codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_raises_value_error_for_unknown_execution_mode', 'run', '4_0', codeflash_loop_index, images=batch, model_version='m', prompt=None, system_prompt=None)

def test_run_remotely_with_no_images_returns_empty_list(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_manager = ModelManager(ModelRegistry({}))
    block = Qwen35VLBlockV1(model_manager=model_manager, api_key=None, step_execution_mode=StepExecutionMode.REMOTE)
    batch = Batch.init(content=[], indices=[])
    called = {'count': 0}

    def fake_infer_lmm(self, inference_input, model_id, prompt=None, model_id_in_path=False):
        called['count'] += 1
        return {'response': 'should-not-be-called'}
    monkeypatch.setattr(InferenceHTTPClient, 'infer_lmm', fake_infer_lmm, raising=True)
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_remotely_with_no_images_returns_empty_list', 'run', '6', codeflash_loop_index, images=batch, model_version='m', prompt=None, system_prompt=None)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_run_remotely_uses_hosted_url_and_selects_v0(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(qwen_module, 'WORKFLOWS_REMOTE_API_TARGET', 'hosted', raising=False)
    selected = {'called': False}
    observed = {}

    def fake_select_api_v0(self):
        selected['called'] = True
        return self

    def fake_infer_lmm(self, inference_input, model_id, prompt=None, model_id_in_path=False):
        api_url_attr = f'_{self.__class__.__name__}__api_url'
        observed['api_url'] = getattr(self, api_url_attr)
        observed['prompt'] = prompt
        return {'response': 'hosted-response'}
    monkeypatch.setattr(InferenceHTTPClient, 'select_api_v0', fake_select_api_v0, raising=True)
    monkeypatch.setattr(InferenceHTTPClient, 'infer_lmm', fake_infer_lmm, raising=True)
    model_manager = ModelManager(ModelRegistry({}))
    block = Qwen35VLBlockV1(model_manager=model_manager, api_key='k', step_execution_mode=StepExecutionMode.REMOTE)
    img = WorkflowImageData(parent_metadata=object(), base64_image='b64-hosted')
    batch = Batch.init(content=[img], indices=[(0,)])
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_remotely_uses_hosted_url_and_selects_v0', 'run', '11', codeflash_loop_index, images=batch, model_version='model-hosted', prompt=None, system_prompt=None)
    codeflash_output = codeflash_return_value
    out = codeflash_output

def test_run_remotely_handles_large_batch_of_images(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    model_manager = ModelManager(ModelRegistry({}))
    block = Qwen35VLBlockV1(model_manager=model_manager, api_key='large-key', step_execution_mode=StepExecutionMode.REMOTE)
    n = 1000
    contents = [WorkflowImageData(parent_metadata=object(), base64_image=f'img-base64-{i}') for i in range(n)]
    indices = [(i,) for i in range(n)]
    batch = Batch.init(content=contents, indices=indices)

    def fake_infer_lmm(self, inference_input, model_id, prompt=None, model_id_in_path=False):
        return {'response': f'response-for-{inference_input}'}
    monkeypatch.setattr(InferenceHTTPClient, 'infer_lmm', fake_infer_lmm, raising=True)
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_remotely_handles_large_batch_of_images', 'run', '8', codeflash_loop_index, images=batch, model_version='large-model', prompt=None, system_prompt=None)
    codeflash_output = codeflash_return_value
    results = codeflash_output