from __future__ import annotations

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

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(
    codeflash_wrapped: Callable[..., Any],
    codeflash_test_module_name: str,
    codeflash_test_class_name: Optional[str],
    codeflash_test_name: str,
    codeflash_function_name: str,
    codeflash_line_id: str,
    codeflash_loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> str:
    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}"
    exception = None
    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}######$!")
    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_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."
    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,
        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,
    )
    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."
    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,
        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,
    )
    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."
    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,
        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,
    )
    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."
    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,
        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,
    )
    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."
    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,
        roboflow_api_key="rf_key:roboflow",openai_api_key="sk-empty",prompt=[],gpt_model_version="gpt-4v",max_tokens=10,temperature=0.0,
    )
    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."
    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,
        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,
    )
    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)."
    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,
        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,
    )
    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."
    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,
        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,
    )
    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."
    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,
        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,
    )
    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."
    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,
        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,
    )
    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)]
    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,
        roboflow_api_key="rf_key:roboflow",openai_api_key="sk-largeprompt",prompt=prompt,gpt_model_version="gpt-4v",max_tokens=100,temperature=0.5,
    )
    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:
        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,
            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,
        )


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:
        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,
            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,
        )


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:
        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,
            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,
        )


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:
        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,
            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,
        )


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)]
    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,
        roboflow_api_key="rf_key:roboflow",openai_api_key="sk-largelist",prompt=prompt,gpt_model_version="gpt-4v",max_tokens=500,temperature=0.5,
    )
    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."
    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,
        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,
    )
    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)]
    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,
        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,
    )
    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"}]
    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,
        roboflow_api_key="rf_key:roboflow",openai_api_key="sk-repeat",prompt=prompt,gpt_model_version="gpt-4v",max_tokens=10,temperature=0.5,
    )
    results = codeflash_return_value
