from __future__ import annotations

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

import dill as pickle
import pytest
from inference.usage_tracking import collector as collector_mod
from inference.usage_tracking.collector import UsageCollector


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 test_non_gcp_serverless_returns_raw_positive(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(collector_mod, 'GCP_SERVERLESS', False)
    monkeypatch.setattr(collector_mod, 'apply_duration_minimum', None)
    t1 = 1.0
    t2 = 2.5
    _call__bound__arguments = inspect.signature(UsageCollector._compute_execution_duration).bind(t1, t2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_0', None, 'test_non_gcp_serverless_returns_raw_positive', '_compute_execution_duration', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_non_gcp_serverless_returns_raw_negative(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(collector_mod, 'GCP_SERVERLESS', False)
    monkeypatch.setattr(collector_mod, 'apply_duration_minimum', None)
    t1 = 5.0
    t2 = 3.0
    _call__bound__arguments = inspect.signature(UsageCollector._compute_execution_duration).bind(t1, t2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_0', None, 'test_non_gcp_serverless_returns_raw_negative', '_compute_execution_duration', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_gcp_serverless_with_no_apply_object_applies_minimum_for_small_raw(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(collector_mod, 'GCP_SERVERLESS', True)
    monkeypatch.setattr(collector_mod, 'apply_duration_minimum', None)
    t1 = 1.0
    t2 = 1.05
    _call__bound__arguments = inspect.signature(UsageCollector._compute_execution_duration).bind(t1, t2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_0', None, 'test_gcp_serverless_with_no_apply_object_applies_minimum_for_small_raw', '_compute_execution_duration', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_gcp_serverless_with_no_apply_object_returns_raw_if_larger_than_min(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(collector_mod, 'GCP_SERVERLESS', True)
    monkeypatch.setattr(collector_mod, 'apply_duration_minimum', None)
    t1 = 1.0
    t2 = 1.5
    _call__bound__arguments = inspect.signature(UsageCollector._compute_execution_duration).bind(t1, t2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_0', None, 'test_gcp_serverless_with_no_apply_object_returns_raw_if_larger_than_min', '_compute_execution_duration', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_apply_duration_minimum_get_true_enforces_minimum(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(collector_mod, 'GCP_SERVERLESS', True)

    class ApplyObjTrue:

        def get(self):
            return True
    monkeypatch.setattr(collector_mod, 'apply_duration_minimum', ApplyObjTrue())
    t1 = 10.0
    t2 = 10.00001
    _call__bound__arguments = inspect.signature(UsageCollector._compute_execution_duration).bind(t1, t2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_0', None, 'test_apply_duration_minimum_get_true_enforces_minimum', '_compute_execution_duration', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_apply_duration_minimum_get_false_returns_raw(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(collector_mod, 'GCP_SERVERLESS', True)

    class ApplyObjFalse:

        def get(self):
            return False
    monkeypatch.setattr(collector_mod, 'apply_duration_minimum', ApplyObjFalse())
    t1 = 3.5
    t2 = 3.51
    _call__bound__arguments = inspect.signature(UsageCollector._compute_execution_duration).bind(t1, t2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_0', None, 'test_apply_duration_minimum_get_false_returns_raw', '_compute_execution_duration', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_apply_duration_minimum_get_raises_lookuperror_uses_minimum(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(collector_mod, 'GCP_SERVERLESS', True)

    class ApplyObjRaises:

        def get(self):
            raise LookupError('config missing')
    monkeypatch.setattr(collector_mod, 'apply_duration_minimum', ApplyObjRaises())
    t1 = 0.0
    t2 = 0.02
    _call__bound__arguments = inspect.signature(UsageCollector._compute_execution_duration).bind(t1, t2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_0', None, 'test_apply_duration_minimum_get_raises_lookuperror_uses_minimum', '_compute_execution_duration', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_serverless_minimum_enforced_on_many_values(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(collector_mod, 'GCP_SERVERLESS', True)

    class ApplyObjTrue:

        def get(self):
            return True
    monkeypatch.setattr(collector_mod, 'apply_duration_minimum', ApplyObjTrue())
    inputs = []
    expected = []
    for i in range(1000):
        if i % 2 == 0:
            t1 = float(i)
            t2 = t1 + 0.001
            inputs.append((t1, t2))
            expected.append(0.1)
        else:
            t1 = float(i)
            t2 = t1 + 0.5
            inputs.append((t1, t2))
            expected.append(0.5)
    results = [UsageCollector._compute_execution_duration(a, b) for a, b in inputs]
    for idx, (res, exp) in enumerate(zip(results, expected)):
        pass