from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path
from unittest import mock

import dill as pickle
import numpy as np
import pytest
import torch
from inference.core.workflows.execution_engine.entities.base import (
    ImageParentMetadata, WorkflowImageData)
from inference.models.depth_anything_v3.depth_anything_v3_inference_models import \
    InferenceModelsDepthAnythingV3Adapter


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}######!')
    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_returning(tensor_to_return, record_inputs_list=None):
    """
    Helper that builds a callable object which, when called with inputs, records the inputs
    (if record_inputs_list is provided) and returns a list-like where index 0 is the
    provided torch tensor. This mirrors the minimal runtime contract used by predict.
    We return a plain function (not a class) to avoid introducing new domain classes.
    """

    def fake_model_callable(received_inputs):
        if record_inputs_list is not None:
            record_inputs_list.append(received_inputs)
        return [tensor_to_return]
    return fake_model_callable

def test_basic_prediction_valid_output(monkeypatch):
    np.random.seed(42)
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    depth = torch.tensor([[0.0, 1.0], [2.0, 3.0]], dtype=torch.float32)
    target_path = 'inference.models.depth_anything_v3.depth_anything_v3_inference_models.AutoModel.from_pretrained'
    fake_model = _make_fake_model_returning(depth)
    monkeypatch.setattr(target_path, lambda *args, **kwargs: fake_model)
    adapter = InferenceModelsDepthAnythingV3Adapter(model_id='fake-model-id')
    dummy_input = np.zeros((2, 2), dtype=np.float32)
    _call__bound__arguments = inspect.signature(adapter.predict).bind(dummy_input)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_basic_prediction_valid_output', 'predict', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result_tuple = codeflash_output
    result = result_tuple[0]
    normalized = result['normalized_depth']
    image_obj = result['image']
    img_array = image_obj.numpy_image
    parent_meta = image_obj.parent_metadata

def test_predict_raises_on_constant_depth_map(monkeypatch):
    np.random.seed(42)
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    const_depth = torch.ones((3, 3), dtype=torch.float32) * 5.0
    target_path = 'inference.models.depth_anything_v3.depth_anything_v3_inference_models.AutoModel.from_pretrained'
    fake_model = _make_fake_model_returning(const_depth)
    monkeypatch.setattr(target_path, lambda *args, **kwargs: fake_model)
    adapter = InferenceModelsDepthAnythingV3Adapter(model_id='fake-const')
    dummy_input = np.zeros((3, 3), dtype=np.float32)
    with pytest.raises(ValueError):
        _call__bound__arguments = inspect.signature(adapter.predict).bind(dummy_input)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_predict_raises_on_constant_depth_map', 'predict', '6_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_predict_normalization_with_negative_values_and_kwargs(monkeypatch):
    np.random.seed(42)
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    depth_values = torch.tensor([[-10.0, 0.0], [50.0, 90.0]], dtype=torch.float32)
    target_path = 'inference.models.depth_anything_v3.depth_anything_v3_inference_models.AutoModel.from_pretrained'
    fake_model = _make_fake_model_returning(depth_values)
    monkeypatch.setattr(target_path, lambda *args, **kwargs: fake_model)
    adapter = InferenceModelsDepthAnythingV3Adapter(model_id='fake-neg')
    dummy_input = np.ones((2, 2), dtype=np.float32)
    _call__bound__arguments = inspect.signature(adapter.predict).bind(dummy_input, some_unused_kw='value')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_predict_normalization_with_negative_values_and_kwargs', 'predict', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result_tuple = codeflash_output
    normalized = result_tuple[0]['normalized_depth']
    expected_mid = (0.0 - -10.0) / (90.0 - -10.0)

def test_inputs_are_forwarded_to_model(monkeypatch):
    np.random.seed(42)
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    recorded = []
    depth = torch.tensor([[0.0, 1.0]], dtype=torch.float32)
    fake_model = _make_fake_model_returning(depth, record_inputs_list=recorded)
    target_path = 'inference.models.depth_anything_v3.depth_anything_v3_inference_models.AutoModel.from_pretrained'
    monkeypatch.setattr(target_path, lambda *args, **kwargs: fake_model)
    adapter = InferenceModelsDepthAnythingV3Adapter(model_id='fake-forward')
    input_array = np.arange(6, dtype=np.float32).reshape((2, 3))
    _call__bound__arguments = inspect.signature(adapter.predict).bind(input_array)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_inputs_are_forwarded_to_model', 'predict', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    recorded_input = recorded[0]

def test_large_scale_prediction_within_constraints(monkeypatch):
    np.random.seed(42)
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    height, width = (25, 40)
    ramp = np.linspace(0.0, 1.0, num=height * width, dtype=np.float32).reshape((height, width))
    depth_tensor = torch.from_numpy(ramp * 100.0).to(torch.float32)
    target_path = 'inference.models.depth_anything_v3.depth_anything_v3_inference_models.AutoModel.from_pretrained'
    fake_model = _make_fake_model_returning(depth_tensor)
    monkeypatch.setattr(target_path, lambda *args, **kwargs: fake_model)
    adapter = InferenceModelsDepthAnythingV3Adapter(model_id='fake-large')
    input_array = np.zeros((height, width), dtype=np.float32)
    _call__bound__arguments = inspect.signature(adapter.predict).bind(input_array)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_large_scale_prediction_within_constraints', 'predict', '8', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result_tuple = codeflash_output
    normalized = result_tuple[0]['normalized_depth']
    image_obj = result_tuple[0]['image']
    img = image_obj.numpy_image