from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

import cv2
import dill as pickle
import numpy as np
import pytest
from inference.models.moondream2.moondream2_inference_models import \
    InferenceModelsMoondream2Adapter
from PIL import Image


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_fake_model(return_func):
    """
    Create a simple object with a caption(callable) method.
    - return_func: callable that accepts the image arg and returns the first-element-of-caption value.
    The .caption method will return a list/tuple where the desired caption is at index 0,
    mirroring the real implementation's expected return shape.
    """
    recorded = {'last_arg': None, 'call_count': 0}

    def caption(arg):
        recorded['last_arg'] = arg
        recorded['call_count'] += 1
        return [return_func(arg), 'extra_metadata']
    fake = SimpleNamespace(caption=caption, _record=recorded)
    return fake

def test_caption_with_pil_rgb_image_basic():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Basic functionality:\n    - Given a small RGB PIL Image, the adapter should convert it to a BGR numpy array\n      and pass it to the underlying model.caption method, returning the first element.\n    - We validate both the returned value and that the pixel channels have been swapped (RGB->BGR).\n    '
    pil_img = Image.new('RGB', (2, 2), color=(255, 0, 0))

    def return_func(arg):
        top_left = arg[0, 0].tolist()
        return 'caption_for_pil_rgb'
    fake_model = _make_fake_model(return_func)
    with patch('inference.models.moondream2.moondream2_inference_models.AutoModel.from_pretrained', return_value=fake_model):
        adapter = InferenceModelsMoondream2Adapter(model_id='dummy-model-id', api_key='dummy')
        _call__bound__arguments = inspect.signature(adapter.caption).bind(pil_img)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(adapter.caption, 'tests.inference.unit_tests.test_caption__unit_test_0', None, 'test_caption_with_pil_rgb_image_basic', 'caption', '4_1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        out = codeflash_output

def test_caption_with_numpy_bgr_array_identity_preserved():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Basic functionality for numpy input:\n    - When a numpy array is passed, the adapter should NOT convert it and should pass the exact object\n      through to model.caption (i.e., identity should be preserved).\n    - The adapter should return the first element from the underlying caption result.\n    '
    bgr_arr = np.zeros((2, 2, 3), dtype=np.uint8)
    bgr_arr[0, 0] = [10, 20, 30]

    def return_func(arg):
        return 'caption_for_numpy_bgr'
    fake_model = _make_fake_model(return_func)
    with patch('inference.models.moondream2.moondream2_inference_models.AutoModel.from_pretrained', return_value=fake_model):
        adapter = InferenceModelsMoondream2Adapter(model_id='dummy-model-id', api_key='dummy')
        _call__bound__arguments = inspect.signature(adapter.caption).bind(bgr_arr)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(adapter.caption, 'tests.inference.unit_tests.test_caption__unit_test_0', None, 'test_caption_with_numpy_bgr_array_identity_preserved', 'caption', '5_1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        out = codeflash_output

def test_caption_with_grayscale_pil_image_raises_cv2_error():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "\n    Edge case:\n    - A grayscale ('L') PIL Image will produce a 2D numpy array.\n    - cv2.cvtColor with COLOR_RGB2BGR expects a 3-channel image and should raise cv2.error.\n    - The adapter does not handle this explicitly, so we expect a cv2.error to propagate.\n    "
    gray_img = Image.new('L', (2, 2), color=128)
    fake_model = _make_fake_model(lambda arg: "won't be used")
    with patch('inference.models.moondream2.moondream2_inference_models.AutoModel.from_pretrained', return_value=fake_model):
        adapter = InferenceModelsMoondream2Adapter(model_id='dummy-model-id', api_key='dummy')
        with pytest.raises(cv2.error):
            _call__bound__arguments = inspect.signature(adapter.caption).bind(gray_img)
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(adapter.caption, 'tests.inference.unit_tests.test_caption__unit_test_0', None, 'test_caption_with_grayscale_pil_image_raises_cv2_error', 'caption', '3_1_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_caption_with_none_raises_exception():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Edge case:\n    - Passing None should lead numpy to create an object-dtype array, which then causes cv2.cvtColor to raise.\n    - We assert an exception occurs (cv2.error is expected) to ensure the adapter does not silently accept invalid input.\n    '
    fake_model = _make_fake_model(lambda arg: 'unused')
    with patch('inference.models.moondream2.moondream2_inference_models.AutoModel.from_pretrained', return_value=fake_model):
        adapter = InferenceModelsMoondream2Adapter(model_id='dummy-model-id', api_key='dummy')
        with pytest.raises(Exception):
            _call__bound__arguments = inspect.signature(adapter.caption).bind(None)
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(adapter.caption, 'tests.inference.unit_tests.test_caption__unit_test_0', None, 'test_caption_with_none_raises_exception', 'caption', '2_1_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_scale_repeated_captions_small_images_loop():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Large-scale scenario (scaled within reasonable test limits):\n    - Call adapter.caption repeatedly (100 times) with small images to simulate sustained usage.\n    - Each call uses a tiny image (10x10x3 => 300 elements, well under the 1000-element guideline).\n    - Validate correctness across multiple calls and that the underlying model.caption is invoked each time.\n    '
    calls = []

    def return_func(arg):
        s = int(np.sum(arg))
        caption_text = f'sum_{s}'
        calls.append(caption_text)
        return caption_text
    fake_model = _make_fake_model(return_func)
    with patch('inference.models.moondream2.moondream2_inference_models.AutoModel.from_pretrained', return_value=fake_model):
        adapter = InferenceModelsMoondream2Adapter(model_id='dummy-model-id', api_key='dummy')
        num_calls = 100
        results = []
        for i in range(num_calls):
            color = (i % 256, i * 2 % 256, i * 3 % 256)
            img = Image.new('RGB', (10, 10), color=color)
            _call__bound__arguments = inspect.signature(adapter.caption).bind(img)
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(adapter.caption, 'tests.inference.unit_tests.test_caption__unit_test_0', None, 'test_large_scale_repeated_captions_small_images_loop', 'caption', '4_3_2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
            codeflash_output = codeflash_return_value
            res = codeflash_output
            results.append(res)
        for r, c in zip(results, calls):
            pass