from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path
from typing import List, Tuple

import dill as pickle
import pytest
from inference.core.managers.base import ModelManager
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.http.client 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}######!')
    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

def make_workflow_image_data(base64_value: str) -> WorkflowImageData:
    """
    Helper to construct a WorkflowImageData instance.
    We pass a simple built-in object() as parent_metadata because the run_remotely
    implementation only reads the .base64_image property (which we provide).
    """
    parent_metadata = object()
    return WorkflowImageData(parent_metadata=parent_metadata, base64_image=base64_value)

def make_batch_from_base64_list(base64_list: List[str]) -> Batch:
    """
    Helper to wrap a list of base64 strings into a Batch[WorkflowImageData] instance
    using the real Batch.init classmethod (ensures valid indices).
    """
    content = [make_workflow_image_data(b64) for b64 in base64_list]
    indices: List[Tuple[int, ...]] = [(i,) for i in range(len(content))]
    return Batch.init(content=content, indices=indices)

def make_block(api_key: str=None, step_mode: StepExecutionMode=StepExecutionMode.REMOTE):
    """Create a real Qwen35VLBlockV1 instance with a minimal ModelManager."""
    mm = ModelManager(None)
    return Qwen35VLBlockV1(model_manager=mm, api_key=api_key, step_execution_mode=step_mode)

def test_run_remotely_single_image_default_prompts(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "\n    Basic functionality: for a single image, run_remotely should call infer_lmm and\n    return a list with one parsed_output entry equal to the 'response' value from the client.\n    "
    block = make_block()
    batch = make_batch_from_base64_list(['base64-image-1'])

    def fake_infer_lmm(self, inference_input, model_id, prompt=None, model_id_in_path=False):
        return {'response': f'echo:{inference_input}'}
    monkeypatch.setattr(InferenceHTTPClient, 'infer_lmm', fake_infer_lmm)
    _call__bound__arguments = inspect.signature(block.run_remotely).bind(images=batch, model_version='qwen3', prompt=None, system_prompt=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run_remotely, 'tests.inference.unit_tests.test_run_remotely__unit_test_0', None, 'test_run_remotely_single_image_default_prompts', 'run_remotely', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_run_remotely_explicit_prompts_and_combination(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Ensure that provided prompt and system_prompt are combined and passed to infer_lmm.\n    The block uses \'prompt + "<system_prompt>" + system_prompt\' to create combined_prompt.\n    We\'ll return that combined prompt as the response to verify correctness.\n    '
    block = make_block()
    batch = make_batch_from_base64_list(['img-A'])
    user_prompt = 'What is in the image?'
    system_prompt = 'Be concise.'

    def fake_infer_lmm(self, inference_input, model_id, prompt=None, model_id_in_path=False):
        return {'response': prompt}
    monkeypatch.setattr(InferenceHTTPClient, 'infer_lmm', fake_infer_lmm)
    _call__bound__arguments = inspect.signature(block.run_remotely).bind(images=batch, model_version='qwen3-v1', prompt=user_prompt, system_prompt=system_prompt)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run_remotely, 'tests.inference.unit_tests.test_run_remotely__unit_test_0', None, 'test_run_remotely_explicit_prompts_and_combination', 'run_remotely', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected_combined = user_prompt + '<system_prompt>' + system_prompt

def test_run_remotely_result_without_response_key_returns_full_result(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "\n    If infer_lmm returns a dict without 'response' key, run_remotely should use the whole result dict\n    as the parsed_output.\n    "
    block = make_block()
    batch = make_batch_from_base64_list(['img-edge'])

    def fake_infer_lmm(self, inference_input, model_id, prompt=None, model_id_in_path=False):
        return {'answer': 42, 'in': inference_input, 'model': model_id}
    monkeypatch.setattr(InferenceHTTPClient, 'infer_lmm', fake_infer_lmm)
    _call__bound__arguments = inspect.signature(block.run_remotely).bind(images=batch, model_version='m1', prompt=None, system_prompt=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run_remotely, 'tests.inference.unit_tests.test_run_remotely__unit_test_0', None, 'test_run_remotely_result_without_response_key_returns_full_result', 'run_remotely', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_run_remotely_with_empty_batch_does_not_call_client(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    When an empty Batch is provided, run_remotely should return an empty list and\n    must not call the inference client at all.\n    '
    block = make_block()
    empty_batch = Batch.init(content=[], indices=[])

    def failing_infer(self, inference_input, model_id, prompt=None, model_id_in_path=False):
        raise AssertionError('infer_lmm should not be called for an empty batch')
    monkeypatch.setattr(InferenceHTTPClient, 'infer_lmm', failing_infer)
    _call__bound__arguments = inspect.signature(block.run_remotely).bind(images=empty_batch, model_version='any', prompt=None, system_prompt=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run_remotely, 'tests.inference.unit_tests.test_run_remotely__unit_test_0', None, 'test_run_remotely_with_empty_batch_does_not_call_client', 'run_remotely', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_run_remotely_handles_multiple_images(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Verify that run_remotely iterates over all images in the batch and collects predictions\n    in order. This tests behavior with more than one element.\n    '
    block = make_block()
    values = [f'b64-{i}' for i in range(5)]
    batch = make_batch_from_base64_list(values)

    def fake_infer_lmm(self, inference_input, model_id, prompt=None, model_id_in_path=False):
        return {'response': f'resp-for-{inference_input}'}
    monkeypatch.setattr(InferenceHTTPClient, 'infer_lmm', fake_infer_lmm)
    _call__bound__arguments = inspect.signature(block.run_remotely).bind(images=batch, model_version='mv', prompt=None, system_prompt=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run_remotely, 'tests.inference.unit_tests.test_run_remotely__unit_test_0', None, 'test_run_remotely_handles_multiple_images', 'run_remotely', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_run_remotely_large_batch_of_1000_images(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Large-scale test: create a batch of 1000 images and ensure run_remotely processes all of them.\n    The client is monkeypatched to a lightweight deterministic implementation to keep test fast.\n    '
    block = make_block()
    N = 1000
    base64_list = [f'bulk-{i}' for i in range(N)]
    batch = make_batch_from_base64_list(base64_list)

    def fake_infer_lmm(self, inference_input, model_id, prompt=None, model_id_in_path=False):
        return {'response': inference_input}
    monkeypatch.setattr(InferenceHTTPClient, 'infer_lmm', fake_infer_lmm)
    _call__bound__arguments = inspect.signature(block.run_remotely).bind(images=batch, model_version='m-large', prompt=None, system_prompt=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run_remotely, 'tests.inference.unit_tests.test_run_remotely__unit_test_0', None, 'test_run_remotely_large_batch_of_1000_images', 'run_remotely', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output