from __future__ import annotations

import gc
import inspect
import os
import time
import types
from pathlib import Path
from typing import Any, Callable, List, Optional

import dill as pickle
import pytest
import requests
from inference.core.env import API_BASE_URL
from inference.core.workflows.core_steps.models.foundation.openai.v3 import \
    execute_gpt_4v_request
from openai import OpenAI
from openai._types import NOT_GIVEN


def codeflash_wrap(
    wrapped: Callable[..., Any],
    test_module_name: str,
    test_class_name: str | None,
    test_name: str,
    function_name: str,
    line_id: str,
    loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> Any:
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(
        f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!"
    )
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = 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()
    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"{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{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(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_proxied_request_basic():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test proxied request with valid Roboflow and OpenAI keys and a simple prompt."
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="rf_key:account:123",
        prompt=[{"role": "user", "content": "Hello world"}],
        gpt_model_version="gpt-4v",
        max_tokens=100,
        temperature=0.5,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_proxied_request_basic",
        "execute_gpt_4v_request",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_direct_request_basic():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test direct request with a standard OpenAI key and a simple prompt."
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="sk-abc123",
        prompt=[{"role": "user", "content": "Hi there"}],
        gpt_model_version="gpt-4v",
        max_tokens=50,
        temperature=0.7,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_direct_request_basic",
        "execute_gpt_4v_request",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_proxied_user_key():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test proxied request with a user key prefix."
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="rf_key:user:456",
        prompt=[{"role": "user", "content": "Test user key"}],
        gpt_model_version="gpt-4v",
        max_tokens=10,
        temperature=None,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_proxied_user_key",
        "execute_gpt_4v_request",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_direct_request_no_temperature():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test direct request with temperature set to None."
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="sk-xyz789",
        prompt=[{"role": "user", "content": "No temp"}],
        gpt_model_version="gpt-4v",
        max_tokens=20,
        temperature=None,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_direct_request_no_temperature",
        "execute_gpt_4v_request",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_empty_prompt():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test with an empty prompt list."
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="sk-empty",
        prompt=[],
        gpt_model_version="gpt-4v",
        max_tokens=10,
        temperature=0.0,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_empty_prompt",
        "execute_gpt_4v_request",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_empty_prompt_proxied():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test with an empty prompt list via proxy."
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="rf_key:account:empty",
        prompt=[],
        gpt_model_version="gpt-4v",
        max_tokens=10,
        temperature=0.0,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_empty_prompt_proxied",
        "execute_gpt_4v_request",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_invalid_openai_key_prefix():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test with an OpenAI key that has an unknown prefix (should use direct)."
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="unknown_key:abc",
        prompt=[{"role": "user", "content": "Edge"}],
        gpt_model_version="gpt-4v",
        max_tokens=1,
        temperature=1.0,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_invalid_openai_key_prefix",
        "execute_gpt_4v_request",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_max_tokens_zero():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test with max_tokens set to zero."
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="sk-zerotoken",
        prompt=[{"role": "user", "content": "Zero tokens"}],
        gpt_model_version="gpt-4v",
        max_tokens=0,
        temperature=0.9,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_max_tokens_zero",
        "execute_gpt_4v_request",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_temperature_negative():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test with a negative temperature value."
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="sk-negtemp",
        prompt=[{"role": "user", "content": "Negative temp"}],
        gpt_model_version="gpt-4v",
        max_tokens=5,
        temperature=-1.0,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_temperature_negative",
        "execute_gpt_4v_request",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_temperature_above_one():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test with a temperature value above 1."
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="sk-hightemp",
        prompt=[{"role": "user", "content": "High temp"}],
        gpt_model_version="gpt-4v",
        max_tokens=5,
        temperature=2.5,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_temperature_above_one",
        "execute_gpt_4v_request",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_large_prompt():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test with a prompt containing many messages."
    prompt = [{"role": "user", "content": f"msg {i}"} for i in range(10)]
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="sk-largeprompt",
        prompt=prompt,
        gpt_model_version="gpt-4v",
        max_tokens=100,
        temperature=0.5,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_large_prompt",
        "execute_gpt_4v_request",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_proxied_http_error():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test that HTTP errors from the proxy raise RuntimeError."
    with pytest.raises(RuntimeError) as excinfo:
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow",
            openai_api_key="rf_key:account:fail",
            prompt=[{"role": "user", "content": "Fail"}],
            gpt_model_version="fail-model",
            max_tokens=10,
            temperature=0.5,
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            execute_gpt_4v_request,
            "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
            None,
            "test_proxied_http_error",
            "execute_gpt_4v_request",
            "1_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_proxied_bad_structure():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test that a bad response structure from the proxy raises RuntimeError."
    with pytest.raises(RuntimeError) as excinfo:
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow",
            openai_api_key="rf_key:account:bad",
            prompt=[{"role": "user", "content": "Bad structure"}],
            gpt_model_version="bad-structure",
            max_tokens=10,
            temperature=0.5,
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            execute_gpt_4v_request,
            "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
            None,
            "test_proxied_bad_structure",
            "execute_gpt_4v_request",
            "1_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_direct_openai_api_error():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test that errors from the OpenAI API raise RuntimeError."
    with pytest.raises(RuntimeError) as excinfo:
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow",
            openai_api_key="sk-fail",
            prompt=[{"role": "user", "content": "Fail"}],
            gpt_model_version="fail-model",
            max_tokens=10,
            temperature=0.5,
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            execute_gpt_4v_request,
            "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
            None,
            "test_direct_openai_api_error",
            "execute_gpt_4v_request",
            "1_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_direct_bad_structure():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test that a direct response with missing choices raises AttributeError (caught as RuntimeError)."
    with pytest.raises(RuntimeError) as excinfo:
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow",
            openai_api_key="sk-badstructure",
            prompt=[{"role": "user", "content": "Bad structure"}],
            gpt_model_version="bad-structure",
            max_tokens=10,
            temperature=0.5,
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            execute_gpt_4v_request,
            "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
            None,
            "test_direct_bad_structure",
            "execute_gpt_4v_request",
            "1_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_large_prompt_list():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test with a prompt containing 999 messages (upper bound)."
    prompt = [{"role": "user", "content": f"msg {i}"} for i in range(999)]
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="sk-largelist",
        prompt=prompt,
        gpt_model_version="gpt-4v",
        max_tokens=500,
        temperature=0.5,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_large_prompt_list",
        "execute_gpt_4v_request",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_large_token_count():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test with a very high max_tokens value."
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="sk-largetokens",
        prompt=[{"role": "user", "content": "Lots of tokens"}],
        gpt_model_version="gpt-4v",
        max_tokens=999,
        temperature=0.5,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_large_token_count",
        "execute_gpt_4v_request",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_large_scale_proxied():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test proxied request with a large prompt and high max_tokens."
    prompt = [{"role": "user", "content": f"proxied msg {i}"} for i in range(999)]
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="rf_key:user:large",
        prompt=prompt,
        gpt_model_version="gpt-4v",
        max_tokens=999,
        temperature=1.0,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_large_scale_proxied",
        "execute_gpt_4v_request",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_multiple_requests_consistency():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test that multiple calls with same inputs yield same outputs (determinism)."
    prompt = [{"role": "user", "content": "Repeat"}]
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow",
        openai_api_key="sk-repeat",
        prompt=prompt,
        gpt_model_version="gpt-4v",
        max_tokens=10,
        temperature=0.5,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        execute_gpt_4v_request,
        "tests.inference.unit_tests.test_execute_gpt_4v_request__unit_test_1",
        None,
        "test_multiple_requests_consistency",
        "execute_gpt_4v_request",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    results = codeflash_return_value
