from __future__ import annotations

import gc
import os
import time
import types

import pytest
import torch
from inference.core.entities.responses.inference import (
    InferenceResponseImage, ObjectDetectionInferenceResponse,
    ObjectDetectionPrediction)
from inference.core.models.inference_models_adapters import \
    InferenceModelsObjectDetectionAdapter


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

def make_detection_obj(xyxy_tensor, conf_tensor, class_id_tensor):
    """
    Helper to create a detection-like object whose attributes are tensors.
    We avoid defining a new class and instead use types.SimpleNamespace to hold attributes.
    """
    return types.SimpleNamespace(xyxy=xyxy_tensor, confidence=conf_tensor, class_id=class_id_tensor)

def make_preproc_meta(width: int, height: int):
    """
    Build a minimal preprocessing metadata object with original_size.width and original_size.height.
    Use SimpleNamespace to avoid heavy dependencies while preserving required attributes.
    """
    original_size = types.SimpleNamespace(width=width, height=height)
    return types.SimpleNamespace(original_size=original_size)

def make_adapter_with_stub_model(class_names, post_process_return):
    """
    Construct an adapter instance without calling its __init__.
    Assign class_names and a stubbed _model with post_process that returns
    the provided post_process_return value.
    """
    adapter = object.__new__(InferenceModelsObjectDetectionAdapter)
    adapter.class_names = list(class_names)

    def post_process(predictions, preprocess_return_metadata, **kwargs):
        return post_process_return
    adapter._model = types.SimpleNamespace(post_process=post_process)
    return adapter

def test_basic_single_detection_transformed_coordinates_and_metadata():
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    xyxy = torch.tensor([[10.0, 20.0, 30.0, 60.0]])
    conf = torch.tensor([0.85])
    class_id = torch.tensor([1])
    det_obj = make_detection_obj(xyxy, conf, class_id)
    meta = make_preproc_meta(width=640, height=480)
    adapter = make_adapter_with_stub_model(class_names=['cat', 'dog'], post_process_return=[det_obj])
    codeflash_return_value = codeflash_wrap(adapter.postprocess, 'tests.inference.unit_tests.test_postprocess__unit_test_0', None, 'test_basic_single_detection_transformed_coordinates_and_metadata', 'postprocess', '6', codeflash_loop_index, predictions=[], preprocess_return_metadata=[meta])
    codeflash_output = codeflash_return_value
    responses = codeflash_output
    response = responses[0]
    pred = response.predictions[0]

def test_class_id_out_of_range_becomes_string():
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    xyxy = torch.tensor([[0.0, 0.0, 2.0, 4.0]])
    conf = torch.tensor([0.5])
    class_id = torch.tensor([999])
    det_obj = make_detection_obj(xyxy, conf, class_id)
    meta = make_preproc_meta(width=10, height=20)
    adapter = make_adapter_with_stub_model(class_names=['only_one_class'], post_process_return=[det_obj])
    codeflash_return_value = codeflash_wrap(adapter.postprocess, 'tests.inference.unit_tests.test_postprocess__unit_test_0', None, 'test_class_id_out_of_range_becomes_string', 'postprocess', '6', codeflash_loop_index, predictions=[], preprocess_return_metadata=[meta])
    codeflash_output = codeflash_return_value
    responses = codeflash_output
    response = responses[0]
    pred = response.predictions[0]

def test_class_filter_excludes_predictions():
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    xyxy = torch.tensor([[1.0, 1.0, 5.0, 5.0]])
    conf = torch.tensor([0.99])
    class_id = torch.tensor([0])
    det_obj = make_detection_obj(xyxy, conf, class_id)
    meta = make_preproc_meta(width=50, height=50)
    adapter = make_adapter_with_stub_model(class_names=['apple', 'banana'], post_process_return=[det_obj])
    codeflash_return_value = codeflash_wrap(adapter.postprocess, 'tests.inference.unit_tests.test_postprocess__unit_test_0', None, 'test_class_filter_excludes_predictions', 'postprocess', '6', codeflash_loop_index, predictions=[], preprocess_return_metadata=[meta], class_filter=['banana'])
    codeflash_output = codeflash_return_value
    responses = codeflash_output
    response = responses[0]

def test_multiple_images_and_various_detection_counts():
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    xyxy_empty = torch.empty((0, 4))
    conf_empty = torch.empty((0,))
    class_id_empty = torch.empty((0,))
    det_empty = make_detection_obj(xyxy_empty, conf_empty, class_id_empty)
    xyxy_second = torch.tensor([[0.0, 0.0, 4.0, 4.0], [1.0, 2.0, 3.0, 6.0]])
    conf_second = torch.tensor([0.6, 0.7])
    class_id_second = torch.tensor([0, 1])
    det_second = make_detection_obj(xyxy_second, conf_second, class_id_second)
    adapter = make_adapter_with_stub_model(class_names=['a', 'b'], post_process_return=[det_empty, det_second])
    meta1 = make_preproc_meta(width=100, height=100)
    meta2 = make_preproc_meta(width=200, height=300)
    codeflash_return_value = codeflash_wrap(adapter.postprocess, 'tests.inference.unit_tests.test_postprocess__unit_test_0', None, 'test_multiple_images_and_various_detection_counts', 'postprocess', '11', codeflash_loop_index, predictions=[], preprocess_return_metadata=[meta1, meta2])
    codeflash_output = codeflash_return_value
    responses = codeflash_output
    p0 = responses[1].predictions[0]
    p1 = responses[1].predictions[1]

def test_large_scale_many_detections_per_image():
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 200
    x1 = torch.arange(0, n, dtype=torch.float32)
    y1 = torch.arange(0, n, dtype=torch.float32) * 2.0
    x2 = x1 + 1.0
    y2 = y1 + 2.0
    xyxy = torch.stack([x1, y1, x2, y2], dim=1)
    conf = torch.linspace(0.0, 1.0, steps=n)
    class_ids = (torch.arange(n) % 2).to(torch.int64)
    det_obj = make_detection_obj(xyxy, conf, class_ids)
    meta = make_preproc_meta(width=1024, height=768)
    adapter = make_adapter_with_stub_model(class_names=['left', 'right'], post_process_return=[det_obj])
    codeflash_return_value = codeflash_wrap(adapter.postprocess, 'tests.inference.unit_tests.test_postprocess__unit_test_0', None, 'test_large_scale_many_detections_per_image', 'postprocess', '11', codeflash_loop_index, predictions=[], preprocess_return_metadata=[meta])
    codeflash_output = codeflash_return_value
    responses = codeflash_output
    response = responses[0]
    first = response.predictions[0]
    middle = response.predictions[n // 2]
    last = response.predictions[-1]

def test_empty_inputs_return_empty_responses():
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    adapter = make_adapter_with_stub_model(class_names=[], post_process_return=[])
    codeflash_return_value = codeflash_wrap(adapter.postprocess, 'tests.inference.unit_tests.test_postprocess__unit_test_0', None, 'test_empty_inputs_return_empty_responses', 'postprocess', '1', codeflash_loop_index, predictions=[], preprocess_return_metadata=[])
    codeflash_output = codeflash_return_value
    responses = codeflash_output