from __future__ import annotations

import gc
import os
import time

import numpy as np
import pytest
from inference.core.models.base import BaseInference


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_infer_basic_numpy_roundtrip():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Basic functionality: numpy input flows through preprocess->predict->postprocess\n    and the final return value matches what postprocess returns.'

    class SimpleInference(BaseInference):

        def preprocess(self, image, **kwargs):
            return (image, {'info': 'ok'})

        def predict(self, img_in, **kwargs):
            return (img_in * 2,)

        def postprocess(self, predictions, preprocess_return_metadata, **kwargs):
            arr = predictions[0]
            total = int(np.sum(arr))
            return {'total': total, 'meta': preprocess_return_metadata}
    model = SimpleInference()
    image = np.array([[1, 2], [3, 4]], dtype=np.int64)
    codeflash_return_value = codeflash_wrap(model.infer, 'tests.inference.unit_tests.test_infer__unit_test_0', None, 'test_infer_basic_numpy_roundtrip', 'infer', '4', codeflash_loop_index, image)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_infer_forwards_kwargs_to_all_stages():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Ensure that kwargs supplied to infer are forwarded to preprocess, predict and postprocess.'

    class KWArgsEchoInference(BaseInference):

        def __init__(self):
            self.record = []

        def preprocess(self, image, **kwargs):
            self.record.append(('pre', dict(kwargs)))
            return (np.array([0]), {'stage': 'pre'})

        def predict(self, img_in, **kwargs):
            self.record.append(('pred', dict(kwargs)))
            return (np.array([1]),)

        def postprocess(self, predictions, preprocess_return_metadata, **kwargs):
            self.record.append(('post', dict(kwargs)))
            return {'recorded': self.record, 'meta': preprocess_return_metadata}
    model = KWArgsEchoInference()
    codeflash_return_value = codeflash_wrap(model.infer, 'tests.inference.unit_tests.test_infer__unit_test_0', None, 'test_infer_forwards_kwargs_to_all_stages', 'infer', '3', codeflash_loop_index, np.array([7]), user='alice', trace_id=123)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_infer_handles_preprocessed_without_shape_attribute():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "If preprocess returns an object without a 'shape' attribute (e.g., a Python list),\n    BaseInference should proceed without raising an AttributeError. The data should be\n    passed through to predict/postprocess as-is (subject to subclass implementations)."

    class ListFriendlyInference(BaseInference):

        def preprocess(self, image, **kwargs):
            return ([1, 2, 3], {'orig_type': type(image).__name__})

        def predict(self, img_in, **kwargs):
            return (np.array(img_in) + 10,)

        def postprocess(self, predictions, preprocess_return_metadata, **kwargs):
            return {'sum': int(np.sum(predictions[0])), 'meta': preprocess_return_metadata}
    model = ListFriendlyInference()
    codeflash_return_value = codeflash_wrap(model.infer, 'tests.inference.unit_tests.test_infer__unit_test_0', None, 'test_infer_handles_preprocessed_without_shape_attribute', 'infer', '3', codeflash_loop_index, 'a string path or other non-array input')
    codeflash_output = codeflash_return_value
    out = codeflash_output

def test_infer_propagates_exceptions_from_predict():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'If predict raises an exception, infer should not swallow it.'

    class BrokenPredictInference(BaseInference):

        def preprocess(self, image, **kwargs):
            return (np.array([0]), {})

        def predict(self, img_in, **kwargs):
            raise ValueError('predict failed')

        def postprocess(self, predictions, preprocess_return_metadata, **kwargs):
            return 'should not reach here'
    model = BrokenPredictInference()
    with pytest.raises(ValueError, match='predict failed'):
        codeflash_return_value = codeflash_wrap(model.infer, 'tests.inference.unit_tests.test_infer__unit_test_0', None, 'test_infer_propagates_exceptions_from_predict', 'infer', '3_0', codeflash_loop_index, np.array([1, 2, 3]))

def test_large_array_prediction_scaling():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test behavior with larger arrays (close to the requested scale, but kept reasonable).\n    This ensures the infer pipeline handles arrays of size ~1000 elements.'

    class LargeArrayInference(BaseInference):

        def preprocess(self, image, **kwargs):
            arr = np.asarray(image, dtype=np.int32)
            return (arr, {'length': arr.size})

        def predict(self, img_in, **kwargs):
            return (img_in + 1, img_in * 2)

        def postprocess(self, predictions, preprocess_return_metadata, **kwargs):
            a, b = predictions
            return {'mean_sum': float(np.mean(a) + np.mean(b)), 'meta': preprocess_return_metadata}
    size = 1000
    input_array = np.arange(size, dtype=np.int32)
    model = LargeArrayInference()
    codeflash_return_value = codeflash_wrap(model.infer, 'tests.inference.unit_tests.test_infer__unit_test_0', None, 'test_large_array_prediction_scaling', 'infer', '5', codeflash_loop_index, input_array)
    codeflash_output = codeflash_return_value
    out = codeflash_output

def test_infer_stable_under_many_iterations_performance_smoke():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Run infer many times to exercise repeated usage. This is a smoke test for stability\n    under repeated calls (up to 1000 iterations). Keep per-iteration work small so the test\n    stays fast but still meaningful.'

    class CheapInference(BaseInference):

        def preprocess(self, image, **kwargs):
            return (np.array([1], dtype=np.int32), {})

        def predict(self, img_in, **kwargs):
            return (img_in + 2,)

        def postprocess(self, predictions, preprocess_return_metadata, **kwargs):
            return int(predictions[0][0])
    model = CheapInference()
    iterations = 1000
    for i in range(iterations):
        codeflash_return_value = codeflash_wrap(model.infer, 'tests.inference.unit_tests.test_infer__unit_test_0', None, 'test_infer_stable_under_many_iterations_performance_smoke', 'infer', '4_0', codeflash_loop_index, None)
        codeflash_output = codeflash_return_value