from __future__ import annotations

import gc
import os
import time

import inference.core.workflows.core_steps.composition.workflow_caller.v1 as workflow_caller_v1
import pytest
from inference.core.exceptions import (MalformedWorkflowResponseError,
                                       RetryRequestError,
                                       RoboflowAPIForbiddenError,
                                       RoboflowAPINotAuthorizedError,
                                       RoboflowAPIUnsuccessfulRequestError)
from inference.core.workflows.core_steps.composition.workflow_caller.v1 import \
    _fetch_workflow_spec_for_validation


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_returns_spec_on_success(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    recorded = {}

    def fake_get_workflow_specification(api_key, workspace_id, workflow_id, workflow_version_id=None):
        recorded['args'] = (api_key, workspace_id, workflow_id, workflow_version_id)
        return {'name': 'test-workflow', 'specification': {'steps': []}}
    monkeypatch.setattr(workflow_caller_v1, 'get_workflow_specification', fake_get_workflow_specification)
    codeflash_return_value = codeflash_wrap(_fetch_workflow_spec_for_validation, 'tests.inference.unit_tests.test__fetch_workflow_spec_for_validation__unit_test_0', None, 'test_returns_spec_on_success', '_fetch_workflow_spec_for_validation', '3', codeflash_loop_index, api_key='API_KEY_123', workspace_name='my-workspace', workflow_id='my-workflow-id', workflow_version_id='v1')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_returns_none_on_generic_exception(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

    def fake_raises(api_key, workspace_id, workflow_id, workflow_version_id=None):
        raise Exception('something went wrong')
    monkeypatch.setattr(workflow_caller_v1, 'get_workflow_specification', fake_raises)
    codeflash_return_value = codeflash_wrap(_fetch_workflow_spec_for_validation, 'tests.inference.unit_tests.test__fetch_workflow_spec_for_validation__unit_test_0', None, 'test_returns_none_on_generic_exception', '_fetch_workflow_spec_for_validation', '2', codeflash_loop_index, api_key='key', workspace_name='ws', workflow_id='id-that-fails')
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_returns_none_for_specific_api_exceptions(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    exception_instances = [MalformedWorkflowResponseError('malformed'), RetryRequestError('retry', Exception('inner')), RoboflowAPIUnsuccessfulRequestError('unsuccessful'), RoboflowAPIForbiddenError('forbidden'), RoboflowAPINotAuthorizedError('not authorized')]
    for exc in exception_instances:

        def make_raiser(e):

            def raiser(api_key, workspace_id, workflow_id, workflow_version_id=None):
                raise e
            return raiser
        monkeypatch.setattr(workflow_caller_v1, 'get_workflow_specification', make_raiser(exc))
        codeflash_return_value = codeflash_wrap(_fetch_workflow_spec_for_validation, 'tests.inference.unit_tests.test__fetch_workflow_spec_for_validation__unit_test_0', None, 'test_returns_none_for_specific_api_exceptions', '_fetch_workflow_spec_for_validation', '1_2', codeflash_loop_index, api_key='maybe-none', workspace_name='workspace', workflow_id='workflow', workflow_version_id=None)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_propagates_base_exceptions(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

    def raiser(api_key, workspace_id, workflow_id, workflow_version_id=None):
        raise KeyboardInterrupt('boom!')
    monkeypatch.setattr(workflow_caller_v1, 'get_workflow_specification', raiser)
    with pytest.raises(KeyboardInterrupt):
        codeflash_return_value = codeflash_wrap(_fetch_workflow_spec_for_validation, 'tests.inference.unit_tests.test__fetch_workflow_spec_for_validation__unit_test_0', None, 'test_propagates_base_exceptions', '_fetch_workflow_spec_for_validation', '2_0', codeflash_loop_index, api_key='x', workspace_name='ws', workflow_id='id')

def test_parameter_passthrough_with_none_values(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    recorded = {}

    def fake(api_key, workspace_id, workflow_id, workflow_version_id=None):
        recorded['args'] = (api_key, workspace_id, workflow_id, workflow_version_id)
        return {'ok': True}
    monkeypatch.setattr(workflow_caller_v1, 'get_workflow_specification', fake)
    codeflash_return_value = codeflash_wrap(_fetch_workflow_spec_for_validation, 'tests.inference.unit_tests.test__fetch_workflow_spec_for_validation__unit_test_0', None, 'test_parameter_passthrough_with_none_values', '_fetch_workflow_spec_for_validation', '3', codeflash_loop_index, api_key=None, workspace_name='local-ws', workflow_id='wf', workflow_version_id=None)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_many_calls(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

    def fake(api_key, workspace_id, workflow_id, workflow_version_id=None):
        return {'returned_workflow_id': workflow_id}
    monkeypatch.setattr(workflow_caller_v1, 'get_workflow_specification', fake)
    for i in range(1000):
        wid = f'workflow-{i}'
        codeflash_return_value = codeflash_wrap(_fetch_workflow_spec_for_validation, 'tests.inference.unit_tests.test__fetch_workflow_spec_for_validation__unit_test_0', None, 'test_large_scale_many_calls', '_fetch_workflow_spec_for_validation', '2_1', codeflash_loop_index, api_key=f'key-{i}', workspace_name=f'ws-{i % 5}', workflow_id=wid)
        codeflash_output = codeflash_return_value
        result = codeflash_output