from __future__ import annotations

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

import dill as pickle
import pytest
from inference.core.entities.responses.workflows import WorkflowErrorResponse
from inference.core.exceptions import (ContentTypeInvalid, ContentTypeMissing,
                                       CreditsExceededError,
                                       InferenceModelNotFound,
                                       InputImageLoadError,
                                       InvalidMaskDecodeArgument,
                                       InvalidModelIDError, MissingApiKeyError,
                                       WorkspaceStreamQuotaError)
from inference.core.interfaces.http.error_handlers import with_route_exceptions
from inference.core.interfaces.stream_manager.api.errors import \
    ProcessesManagerInvalidPayload
from inference.core.workflows.errors import (ClientCausedStepExecutionError,
                                             StepExecutionError,
                                             WorkflowSyntaxError)
from inference_models.errors import ModelInputError


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 _extract_json_content(resp):
    """
    Try several common attributes on Response objects to retrieve JSON content.
    This helper attempts to be permissive since the test environment provides
    a lightweight replacement of Starlette's JSONResponse.
    """
    if hasattr(resp, 'body'):
        body = getattr(resp, 'body')
        if body is None:
            return None
        if isinstance(body, bytes):
            return json.loads(body.decode('utf-8'))
        if isinstance(body, str):
            return json.loads(body)
    if hasattr(resp, 'render'):
        try:
            rendered = resp.render()
        except TypeError:
            rendered = None
        except Exception:
            rendered = None
        if isinstance(rendered, bytes):
            return json.loads(rendered.decode('utf-8'))
    if hasattr(resp, 'content'):
        content = getattr(resp, 'content')
        if isinstance(content, (dict, list)):
            return content
        if isinstance(content, (str, bytes)):
            return json.loads(content)
    for candidate in ('_content', 'raw_body', 'body_bytes', 'data'):
        if hasattr(resp, candidate):
            val = getattr(resp, candidate)
            if isinstance(val, (dict, list)):
                return val
            if isinstance(val, bytes):
                return json.loads(val.decode('utf-8'))
            if isinstance(val, str):
                return json.loads(val)
    raise AssertionError('Unable to extract JSON content from response; available attributes: %s' % (sorted((k for k in dir(resp) if not k.startswith('__'))),))

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

    def route_ok(x, y=1):
        return {'sum': x + y}
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route_ok)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_decorator_returns_original_result_when_no_exception', 'with_route_exceptions', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped = codeflash_output
    result = wrapped(3, y=4)

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

    def route(_):
        raise ContentTypeInvalid('bad content type')
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_content_type_invalid_returns_400_message', 'with_route_exceptions', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped = codeflash_output
    resp = wrapped(None)
    content = _extract_json_content(resp)

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

    def route(_):
        raise ContentTypeMissing('no content type')
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_content_type_missing_returns_400_message', 'with_route_exceptions', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped = codeflash_output
    resp = wrapped(None)
    content = _extract_json_content(resp)

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

    def route(_):
        raise InputImageLoadError('internal', 'could not decode image')
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_input_image_load_error_includes_public_message', 'with_route_exceptions', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped = codeflash_output
    resp = wrapped(None)
    content = _extract_json_content(resp)

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

    def route(_):
        raise ModelInputError('bad input', 'https://help.example')
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_model_input_error_leads_to_internal_error_due_to_attribute_mismatch', 'with_route_exceptions', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped = codeflash_output
    resp = wrapped(None)
    content = _extract_json_content(resp)

def test_workflow_syntax_error_returns_400_and_includes_blocks_errors():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    blocks = []

    def route(_):
        raise WorkflowSyntaxError(public_message='syntax problem', context='line 1', blocks_errors=blocks)
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_workflow_syntax_error_returns_400_and_includes_blocks_errors', 'with_route_exceptions', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped = codeflash_output
    resp = wrapped(None)
    content = _extract_json_content(resp)

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

    def route(_):
        raise ClientCausedStepExecutionError(block_id='block-42', status_code=418, public_message='client-caused', context='ctx')
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_client_caused_step_execution_error_propagates_status_code_and_block_id', 'with_route_exceptions', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped = codeflash_output
    resp = wrapped(None)
    content = _extract_json_content(resp)
    found_block_ids = [b.get('block_id') for b in content['blocks_errors'] if isinstance(b, dict)]

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

    def route(_):
        raise InferenceModelNotFound('temporarily not ready')
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_inference_model_not_found_sets_retry_after_header_and_503_status', 'with_route_exceptions', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped = codeflash_output
    resp = wrapped(None)
    headers = getattr(resp, 'headers', None)
    if headers is None and hasattr(resp, 'raw_headers'):
        headers = dict(getattr(resp, 'raw_headers'))
    if isinstance(headers, dict):
        pass
    else:
        pass

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

    def route(_):
        raise ValueError('boom')
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_general_unhandled_exception_results_in_500_internal_message', 'with_route_exceptions', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped = codeflash_output
    resp = wrapped(None)
    content = _extract_json_content(resp)

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

    def route_credits(_):
        raise CreditsExceededError('no credits')
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route_credits)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_credits_exceeded_and_workspace_quota_errors_return_expected_codes_and_types', 'with_route_exceptions', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped_credits = codeflash_output
    resp_credits = wrapped_credits(None)
    content_credits = _extract_json_content(resp_credits)

    def route_quota(_):
        raise WorkspaceStreamQuotaError('quota reached for workspace')
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route_quota)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_credits_exceeded_and_workspace_quota_errors_return_expected_codes_and_types', 'with_route_exceptions', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped_quota = codeflash_output
    resp_quota = wrapped_quota(None)
    content_quota = _extract_json_content(resp_quota)

def test_large_scale_many_wrapped_routes_handle_exceptions_quickly():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    wrapped_funcs = []
    for i in range(1000):

        def make_route():

            def route(_):
                raise ContentTypeMissing(f'missing-{i}')
            return route
        _call__bound__arguments = inspect.signature(with_route_exceptions).bind(make_route())
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_large_scale_many_wrapped_routes_handle_exceptions_quickly', 'with_route_exceptions', '1_1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    for idx, wf in enumerate(wrapped_funcs):
        resp = wf(None)

def test_processes_manager_invalid_payload_includes_public_and_inner_error_type():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inner = KeyError('k')
    pm_exc = ProcessesManagerInvalidPayload('private', public_message='payload bad', inner_error=inner)

    def route(_):
        raise pm_exc
    _call__bound__arguments = inspect.signature(with_route_exceptions).bind(route)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(with_route_exceptions, 'tests.inference.unit_tests.test_with_route_exceptions__unit_test_0', None, 'test_processes_manager_invalid_payload_includes_public_and_inner_error_type', 'with_route_exceptions', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    wrapped = codeflash_output
    resp = wrapped(None)
    content = _extract_json_content(resp)