from __future__ import annotations

import gc
import importlib
import inspect
import math
import os
import random
import time
from pathlib import Path

import dill as pickle
import inference.core.telemetry as telemetry
import pytest
from inference.core.telemetry import record_inference


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

class CountMetric:

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

    def add(self, amount, labels):
        self.calls.append((amount, labels))

class DurationMetric:

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

    def record(self, duration, labels):
        self.calls.append((duration, labels))

def test_no_metrics_does_nothing():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'When telemetry._metrics is None the function should be a no-op and not raise.'
    original = telemetry._metrics
    telemetry._metrics = None
    try:
        _call__bound__arguments = inspect.signature(telemetry.record_inference).bind('model-x', 0.5)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(telemetry.record_inference, 'tests.inference.unit_tests.test_record_inference__unit_test_0', None, 'test_no_metrics_does_nothing', 'record_inference', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        telemetry._metrics = original

def test_basic_recording_calls_add_and_record_with_expected_labels():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Verify that a single call records exactly one add and one record with proper args.'
    original = telemetry._metrics
    count = CountMetric()
    duration = DurationMetric()
    telemetry._metrics = {'model_infer_count': count, 'model_infer_duration': duration}
    try:
        _call__bound__arguments = inspect.signature(telemetry.record_inference).bind('my-model', 1.234)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(telemetry.record_inference, 'tests.inference.unit_tests.test_record_inference__unit_test_0', None, 'test_basic_recording_calls_add_and_record_with_expected_labels', 'record_inference', '5_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        amt, labels = count.calls[0]
        dur, dlabels = duration.calls[0]
    finally:
        telemetry._metrics = original

def test_multiple_calls_accumulate_calls_and_preserve_order():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Multiple invocations should append calls in order and preserve arguments.'
    original = telemetry._metrics
    count = CountMetric()
    duration = DurationMetric()
    telemetry._metrics = {'model_infer_count': count, 'model_infer_duration': duration}
    try:
        _call__bound__arguments = inspect.signature(telemetry.record_inference).bind('mA', 0.1)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(telemetry.record_inference, 'tests.inference.unit_tests.test_record_inference__unit_test_0', None, 'test_multiple_calls_accumulate_calls_and_preserve_order', 'record_inference', '5_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        _call__bound__arguments = inspect.signature(telemetry.record_inference).bind('mB', 0.2)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(telemetry.record_inference, 'tests.inference.unit_tests.test_record_inference__unit_test_0', None, 'test_multiple_calls_accumulate_calls_and_preserve_order', 'record_inference', '5_3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        _call__bound__arguments = inspect.signature(telemetry.record_inference).bind('mA', 0.3)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(telemetry.record_inference, 'tests.inference.unit_tests.test_record_inference__unit_test_0', None, 'test_multiple_calls_accumulate_calls_and_preserve_order', 'record_inference', '5_6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    finally:
        telemetry._metrics = original

def test_edge_values_for_model_id_and_duration():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test empty, long, special-character model ids and special float durations (0, negative, NaN, inf).'
    original = telemetry._metrics
    count = CountMetric()
    duration = DurationMetric()
    telemetry._metrics = {'model_infer_count': count, 'model_infer_duration': duration}
    try:
        very_long = 'm' * 1000
        special = 'model:üñïçødé/!@# $%^&*()'
        cases = [('', 0.0), (very_long, -1.0), (special, float('nan')), ('normal', float('inf'))]
        for mid, dur in cases:
            _call__bound__arguments = inspect.signature(telemetry.record_inference).bind(mid, dur)
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(telemetry.record_inference, 'tests.inference.unit_tests.test_record_inference__unit_test_0', None, 'test_edge_values_for_model_id_and_duration', 'record_inference', '5_3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        for i, (mid, dur) in enumerate(cases):
            c_amt, c_labels = count.calls[i]
            d_dur, d_labels = duration.calls[i]
            if math.isnan(dur):
                pass
            elif math.isinf(dur):
                pass
            else:
                pass
    finally:
        telemetry._metrics = original

def test_non_string_model_id_is_allowed_and_preserved_in_labels():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Although type hints say str, the function should accept other types and pass them through.'
    original = telemetry._metrics
    count = CountMetric()
    duration = DurationMetric()
    telemetry._metrics = {'model_infer_count': count, 'model_infer_duration': duration}
    try:
        model_id_int = 42
        _call__bound__arguments = inspect.signature(telemetry.record_inference).bind(model_id_int, 2.5)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(telemetry.record_inference, 'tests.inference.unit_tests.test_record_inference__unit_test_0', None, 'test_non_string_model_id_is_allowed_and_preserved_in_labels', 'record_inference', '5_1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    finally:
        telemetry._metrics = original

def test_missing_metric_keys_raises_keyerror():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'If expected keys are missing in _metrics the function should raise KeyError on access.'
    original = telemetry._metrics
    telemetry._metrics = {}
    try:
        with pytest.raises(KeyError):
            _call__bound__arguments = inspect.signature(telemetry.record_inference).bind('m', 0.1)
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(telemetry.record_inference, 'tests.inference.unit_tests.test_record_inference__unit_test_0', None, 'test_missing_metric_keys_raises_keyerror', 'record_inference', '3_0_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    finally:
        telemetry._metrics = original

def test_metric_objects_without_expected_methods_raise_attributeerror():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "If the metric objects don't expose .add or .record attribute access should raise AttributeError."
    original = telemetry._metrics
    telemetry._metrics = {'model_infer_count': 123, 'model_infer_duration': 456}
    try:
        with pytest.raises(AttributeError):
            _call__bound__arguments = inspect.signature(telemetry.record_inference).bind('m', 0.1)
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(telemetry.record_inference, 'tests.inference.unit_tests.test_record_inference__unit_test_0', None, 'test_metric_objects_without_expected_methods_raise_attributeerror', 'record_inference', '3_0_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    finally:
        telemetry._metrics = original

def test_large_scale_many_calls_is_detected_and_counts_match():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Large-scale test: perform 1000 iterations and ensure all calls are captured deterministically.'
    original = telemetry._metrics
    count = CountMetric()
    duration = DurationMetric()
    telemetry._metrics = {'model_infer_count': count, 'model_infer_duration': duration}
    try:
        rng = random.Random(0)
        model_ids = [f'm{i}' for i in range(10)]
        iterations = 1000
        expected_counts = {mid: 0 for mid in model_ids}
        for _ in range(iterations):
            mid = rng.choice(model_ids)
            dur = rng.random() * 5.0
            _call__bound__arguments = inspect.signature(telemetry.record_inference).bind(mid, dur)
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(telemetry.record_inference, 'tests.inference.unit_tests.test_record_inference__unit_test_0', None, 'test_large_scale_many_calls_is_detected_and_counts_match', 'record_inference', '5_4_2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
            expected_counts[mid] += 1
        observed_counts = {}
        for amt, labels in count.calls:
            model_label = labels['model.id']
            observed_counts[model_label] = observed_counts.get(model_label, 0) + 1
    finally:
        telemetry._metrics = original