from __future__ import annotations

import gc
import json
import os
import time
from typing import Any, Dict, List

import pytest
import requests
from inference.core.workflows.core_steps.composition.workflow_caller.v1 import (
    _RESOLVED_WORKFLOW_OUTPUTS, MAX_WORKFLOW_CALL_DEPTH, WILDCARD_KIND,
    WorkflowCallerBlockV1, _deserialize_output_value, _make_cache_key,
    build_workflow_inputs, call_workflow, execute_workflow_request)
from inference.core.workflows.execution_engine.entities.base import \
    WorkflowImageData
from inference.core.workflows.execution_engine.entities.types import Kind


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_make_cache_key_pairs_version_and_id():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    k1 = _make_cache_key('workflow-1', 'v1')
    k2 = _make_cache_key('workflow-2', None)

def test_build_workflow_inputs_with_workflow_image_data_and_plain_values():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    img = object.__new__(WorkflowImageData)
    img._base64_image = 'BASE64-STRING-XYZ'
    img._numpy_image = None
    img._image_reference = None
    inputs = {'plain': 123, 'image': img, 'text': 'hello'}
    resolved = build_workflow_inputs(inputs)

def test_build_workflow_inputs_large_scale_performance():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    large_inputs = {f'k{i}': i for i in range(1000)}
    resolved = build_workflow_inputs(large_inputs)
    for i in range(0, 1000, 100):
        pass

def test_call_workflow_detects_circular_chain_and_reports_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    workflow_id = 'wf-abc'
    call_chain = 'other,' + workflow_id + ',end'
    error, message, result = call_workflow(api_key=None, workspace_name='ws', workflow_id=workflow_id, inputs={}, request_timeout=1, call_chain=call_chain, workflow_version_id=None)

def test_call_workflow_detects_depth_limit_and_reports_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    entries = [f'id{i}' for i in range(MAX_WORKFLOW_CALL_DEPTH)]
    call_chain = ','.join(entries)
    error, message, result = call_workflow(api_key=None, workspace_name='ws', workflow_id='unique-wf', inputs={}, request_timeout=1, call_chain=call_chain, workflow_version_id=None)

def test_execute_workflow_request_success_with_real_requests_response(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    resp = requests.Response()
    resp.status_code = 200
    resp._content = b'{"outputs": [{"out1": 42, "out2": "ok"}]}'

    def fake_post(url, json=None, headers=None, timeout=None):
        return resp
    monkeypatch.setattr('requests.post', fake_post)
    error, message, result = execute_workflow_request(api_key='api-key', workspace_name='ws', workflow_id='wf1', inputs={'a': 1}, request_timeout=5, call_chain='chain', workflow_version_id=None)

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

    def raise_exc(url, json=None, headers=None, timeout=None):
        raise RuntimeError('network failure')
    monkeypatch.setattr('requests.post', raise_exc)
    error, message, result = execute_workflow_request(api_key=None, workspace_name='ws', workflow_id='wf-error', inputs={}, request_timeout=1, call_chain='', workflow_version_id=None)

def test_deserialize_output_value_returns_none_for_none_raw_value():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    out = _deserialize_output_value(output_name='o', raw_value=None, output_kinds=[], deserializers={})

def test_deserialize_output_value_returns_raw_for_wildcard_kind():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    out = _deserialize_output_value(output_name='o', raw_value={'x': 1}, output_kinds=[WILDCARD_KIND, WILDCARD_KIND], deserializers={})

def test_deserialize_output_value_raises_when_all_deserializers_fail():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    k1 = Kind(name='k1')
    k2 = Kind(name='k2')

    def d1(name, value):
        raise ValueError('bad1')

    def d2(name, value):
        raise TypeError('bad2')
    deserializers = {'k1': d1, 'k2': d2}
    with pytest.raises(RuntimeError) as exc:
        _deserialize_output_value(output_name='out_err', raw_value='raw_data', output_kinds=[k1, k2], deserializers=deserializers)
    msg = str(exc.value)

def test_run_raises_when_workspace_cannot_be_resolved(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    block = WorkflowCallerBlockV1(api_key=None)
    with pytest.raises(RuntimeError) as exc:
        codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_raises_when_workspace_cannot_be_resolved', 'run', '1_0', codeflash_loop_index, workflow_id='wf-no-ws', inputs={}, request_timeout=1, input_definitions={}, output_definitions={}, workflow_version_id=None)

def test_run_with_resolved_outputs_uses_cache_and_wildcard_deserialization(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr('inference.core.workflows.core_steps.composition.workflow_caller.v1.get_roboflow_workspace', lambda api_key: 'ws-1')

    def fake_execute_workflow_request(api_key, workspace_name, workflow_id, inputs, request_timeout, call_chain, workflow_version_id=None):
        return (False, 'ok', {'my_output': 999})
    monkeypatch.setattr('inference.core.workflows.core_steps.composition.workflow_caller.v1.execute_workflow_request', fake_execute_workflow_request)
    cache_key = ('target-wf', None)
    _RESOLVED_WORKFLOW_OUTPUTS[cache_key] = {'my_output': [WILDCARD_KIND]}
    block = WorkflowCallerBlockV1(api_key='dummy-api-key')
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_with_resolved_outputs_uses_cache_and_wildcard_deserialization', 'run', '6', codeflash_loop_index, workflow_id='target-wf', inputs={'inp': 1}, request_timeout=2, input_definitions={}, output_definitions={}, workflow_version_id=None)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_run_with_no_output_resolution_returns_generic_result(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr('inference.core.workflows.core_steps.composition.workflow_caller.v1.get_roboflow_workspace', lambda api_key: 'ws-1')

    def fake_execute_workflow_request(api_key, workspace_name, workflow_id, inputs, request_timeout, call_chain, workflow_version_id=None):
        return (False, 'ok', {'a': 1, 'b': 2})
    monkeypatch.setattr('inference.core.workflows.core_steps.composition.workflow_caller.v1.execute_workflow_request', fake_execute_workflow_request)
    _RESOLVED_WORKFLOW_OUTPUTS.pop(('wf-no-cache', None), None)
    block = WorkflowCallerBlockV1(api_key='akey')
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_with_no_output_resolution_returns_generic_result', 'run', '5', codeflash_loop_index, workflow_id='wf-no-cache', inputs={}, request_timeout=1, input_definitions={}, output_definitions={}, workflow_version_id=None)
    codeflash_output = codeflash_return_value
    res = codeflash_output

def test_run_propagates_execution_error_from_call_workflow(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr('inference.core.workflows.core_steps.composition.workflow_caller.v1.get_roboflow_workspace', lambda api_key: 'ws-1')
    monkeypatch.setattr('inference.core.workflows.core_steps.composition.workflow_caller.v1.call_workflow', lambda api_key, workspace_name, workflow_id, inputs, request_timeout, call_chain, workflow_version_id=None: (True, 'remote-failure-message', {}))
    block = WorkflowCallerBlockV1(api_key='akey')
    with pytest.raises(RuntimeError) as excinfo:
        codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_propagates_execution_error_from_call_workflow', 'run', '3_0', codeflash_loop_index, workflow_id='wf-fail', inputs={}, request_timeout=1, input_definitions={}, output_definitions={}, workflow_version_id=None)