from __future__ import annotations

import gc
import os
import time

import pytest
from inference.usage_tracking import collector as collector_module
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}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

def test_basic_non_gcp_returns_raw(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(collector_module, 'GCP_SERVERLESS', False)
    monkeypatch.setattr(collector_module, 'apply_duration_minimum', None)
    t1 = 1.0
    t2 = 2.5
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_1', None, 'test_basic_non_gcp_returns_raw', '_compute_execution_duration', '4', codeflash_loop_index, t1, t2)
    codeflash_output = codeflash_return_value
    t1 = 5.0
    t2 = 3.0
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_1', None, 'test_basic_non_gcp_returns_raw', '_compute_execution_duration', '7', codeflash_loop_index, t1, t2)
    codeflash_output = codeflash_return_value

def test_gcp_no_apply_uses_minimum_when_small(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(collector_module, 'GCP_SERVERLESS', True)
    monkeypatch.setattr(collector_module, 'apply_duration_minimum', None)
    t1 = 10.0
    t2 = 10.05
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_1', None, 'test_gcp_no_apply_uses_minimum_when_small', '_compute_execution_duration', '4', codeflash_loop_index, t1, t2)
    codeflash_output = codeflash_return_value
    t1 = 0.0
    t2 = 0.2
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_1', None, 'test_gcp_no_apply_uses_minimum_when_small', '_compute_execution_duration', '7', codeflash_loop_index, t1, t2)
    codeflash_output = codeflash_return_value

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

    class ApplyTrue:

        def get(self):
            return True
    monkeypatch.setattr(collector_module, 'apply_duration_minimum', ApplyTrue())
    t1 = 2.0
    t2 = 2.03
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_1', None, 'test_gcp_with_apply_get_true_enforces_minimum', '_compute_execution_duration', '5', codeflash_loop_index, t1, t2)
    codeflash_output = codeflash_return_value
    t1 = 2.0
    t2 = 2.5
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_1', None, 'test_gcp_with_apply_get_true_enforces_minimum', '_compute_execution_duration', '8', codeflash_loop_index, t1, t2)
    codeflash_output = codeflash_return_value

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

    class ApplyFalse:

        def get(self):
            return False
    monkeypatch.setattr(collector_module, 'apply_duration_minimum', ApplyFalse())
    t1 = 4.0
    t2 = 4.02
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_1', None, 'test_gcp_with_apply_get_false_returns_raw_even_if_small', '_compute_execution_duration', '5', codeflash_loop_index, t1, t2)
    codeflash_output = codeflash_return_value
    t1 = 4.0
    t2 = 4.3
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_1', None, 'test_gcp_with_apply_get_false_returns_raw_even_if_small', '_compute_execution_duration', '8', codeflash_loop_index, t1, t2)
    codeflash_output = codeflash_return_value

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

    class ApplyRaisesLookupError:

        def get(self):
            raise LookupError('simulated missing config')
    monkeypatch.setattr(collector_module, 'apply_duration_minimum', ApplyRaisesLookupError())
    t1 = 7.0
    t2 = 7.02
    codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_1', None, 'test_apply_get_raising_lookuperror_falls_back_to_minimum', '_compute_execution_duration', '5', codeflash_loop_index, t1, t2)
    codeflash_output = codeflash_return_value

def test_large_scale_all_promoted_when_no_apply_and_gcp(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(collector_module, 'GCP_SERVERLESS', True)
    monkeypatch.setattr(collector_module, 'apply_duration_minimum', None)
    results = []
    for i in range(1000):
        t1 = i * 0.001
        t2 = t1 + i % 5 * 0.02
        codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_1', None, 'test_large_scale_all_promoted_when_no_apply_and_gcp', '_compute_execution_duration', '3_2', codeflash_loop_index, t1, t2)

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

    class ApplyFalse:

        def get(self):
            return False
    monkeypatch.setattr(collector_module, 'apply_duration_minimum', ApplyFalse())
    for i in range(1000):
        t1 = i * 0.0007
        increment = i % 7 * 0.01
        t2 = t1 + increment
        expected = increment
        codeflash_return_value = codeflash_wrap(UsageCollector._compute_execution_duration, 'tests.inference.unit_tests.test__compute_execution_duration__unit_test_1', None, 'test_large_scale_with_apply_get_false_returns_varied_raws', '_compute_execution_duration', '3_4', codeflash_loop_index, t1, t2)
        codeflash_output = codeflash_return_value