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


class DummyResponse:
    """A dummy response object for simulating requests responses."""

    def __init__(self, json_data=None, status_code=200, raise_exc=None, text=""):
        self._json_data = json_data
        self.status_code = status_code
        self._raise_exc = raise_exc
        self.text = text

    def json(self):
        if isinstance(self._json_data, Exception):
            raise self._json_data
        return self._json_data

    def raise_for_status(self):
        if self._raise_exc:
            raise self._raise_exc


class DummyOpenAIChoices:

    def __init__(self, content):
        self.message = types.SimpleNamespace(content=content)


class DummyOpenAIResponse:

    def __init__(self, content):
        self.choices = [DummyOpenAIChoices(content)]


class DummyOpenAIChatCompletions:

    def __init__(self, content):
        self._content = content

    def create(self, model, messages, max_tokens, temperature):
        return DummyOpenAIResponse(self._content)


@pytest.fixture
def patch_requests_post(monkeypatch):
    """Patch requests.post to return a dummy response."""

    def _patch(json_data=None, status_code=200, raise_exc=None, text=""):

        def dummy_post(*args, **kwargs):
            return DummyResponse(
                json_data=json_data,
                status_code=status_code,
                raise_exc=raise_exc,
                text=text,
            )

        monkeypatch.setattr(requests, "post", dummy_post)

    return _patch


@pytest.fixture
def patch_openai_client(monkeypatch):
    """Patch OpenAI to return a dummy client."""

    def _patch(content):

        def dummy_openai_init(self, api_key):
            self.api_key = api_key
            self.chat = types.SimpleNamespace(
                completions=DummyOpenAIChatCompletions(content)
            )

        monkeypatch.setattr(OpenAI, "__init__", dummy_openai_init)
        monkeypatch.setattr(OpenAI, "chat", property(lambda self: self.chat))

    return _patch


def test_proxied_request_basic_success(patch_requests_post):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test basic proxied request with valid keys and prompt."
    expected_content = "Hello from proxy!"
    patch_requests_post(
        json_data={"choices": [{"message": {"content": expected_content}}]}
    )
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow123",
        openai_api_key="rf_key:account:openai123",
        prompt=[{"role": "user", "content": "Say hello"}],
        gpt_model_version="gpt-4-vision-preview",
        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_0",
        None,
        "test_proxied_request_basic_success",
        "execute_gpt_4v_request",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_direct_request_basic_success(patch_openai_client):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test basic direct OpenAI request with valid key and prompt."
    expected_content = "Hello from OpenAI!"
    patch_openai_client(expected_content)
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow123",
        openai_api_key="sk-abc123",
        prompt=[{"role": "user", "content": "Say hello"}],
        gpt_model_version="gpt-4-vision-preview",
        max_tokens=10,
        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_0",
        None,
        "test_direct_request_basic_success",
        "execute_gpt_4v_request",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_proxied_request_temperature_none(patch_requests_post):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test proxied request with temperature=None (should omit temperature in payload)."
    expected_content = "No temperature set."
    patch_requests_post(
        json_data={"choices": [{"message": {"content": expected_content}}]}
    )
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow123",
        openai_api_key="rf_key:user:openai123",
        prompt=[{"role": "user", "content": "What's up?"}],
        gpt_model_version="gpt-4-vision-preview",
        max_tokens=5,
        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_0",
        None,
        "test_proxied_request_temperature_none",
        "execute_gpt_4v_request",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_direct_request_temperature_none(patch_openai_client):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test direct request with temperature=None (should use NOT_GIVEN)."
    expected_content = "No temperature set."
    patch_openai_client(expected_content)
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow123",
        openai_api_key="sk-xyz456",
        prompt=[{"role": "user", "content": "What's up?"}],
        gpt_model_version="gpt-4-vision-preview",
        max_tokens=5,
        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_0",
        None,
        "test_direct_request_temperature_none",
        "execute_gpt_4v_request",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_proxied_request_invalid_response_structure(patch_requests_post):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test proxied request with missing keys in response (should raise RuntimeError)."
    patch_requests_post(json_data={"unexpected": "structure"}, text="bad json")
    with pytest.raises(RuntimeError) as excinfo:
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow123",
            openai_api_key="rf_key:account:openai123",
            prompt=[{"role": "user", "content": "Say hello"}],
            gpt_model_version="gpt-4-vision-preview",
            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_0",
            None,
            "test_proxied_request_invalid_response_structure",
            "execute_gpt_4v_request",
            "2_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_proxied_request_http_error(patch_requests_post):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test proxied request where HTTP error occurs (should raise RuntimeError)."
    patch_requests_post(
        json_data={"choices": [{"message": {"content": "irrelevant"}}]},
        raise_exc=requests.exceptions.HTTPError("Bad request"),
    )
    with pytest.raises(RuntimeError) as excinfo:
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow123",
            openai_api_key="rf_key:account:openai123",
            prompt=[{"role": "user", "content": "Say hello"}],
            gpt_model_version="gpt-4-vision-preview",
            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_0",
            None,
            "test_proxied_request_http_error",
            "execute_gpt_4v_request",
            "2_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_direct_request_openai_exception(monkeypatch):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test direct OpenAI request where OpenAI raises an exception (should raise RuntimeError)."

    class FailingCompletions:

        def create(self, *args, **kwargs):
            raise Exception("OpenAI API is down")

    class FailingClient:

        def __init__(self, api_key):
            self.chat = types.SimpleNamespace(completions=FailingCompletions())

    monkeypatch.setattr(OpenAI, "__init__", lambda self, api_key: None)
    monkeypatch.setattr(
        OpenAI, "chat", property(lambda self: FailingClient("api_key").chat)
    )
    with pytest.raises(RuntimeError) as excinfo:
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow123",
            openai_api_key="sk-fail123",
            prompt=[{"role": "user", "content": "Test"}],
            gpt_model_version="gpt-4-vision-preview",
            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_0",
            None,
            "test_direct_request_openai_exception",
            "execute_gpt_4v_request",
            "5_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_empty_prompt_proxied(patch_requests_post):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test proxied request with empty prompt list."
    expected_content = "Empty prompt handled"
    patch_requests_post(
        json_data={"choices": [{"message": {"content": expected_content}}]}
    )
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow123",
        openai_api_key="rf_key:account:openai123",
        prompt=[],
        gpt_model_version="gpt-4-vision-preview",
        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_0",
        None,
        "test_empty_prompt_proxied",
        "execute_gpt_4v_request",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_empty_prompt_direct(patch_openai_client):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test direct request with empty prompt list."
    expected_content = "Empty prompt handled"
    patch_openai_client(expected_content)
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow123",
        openai_api_key="sk-abc123",
        prompt=[],
        gpt_model_version="gpt-4-vision-preview",
        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_0",
        None,
        "test_empty_prompt_direct",
        "execute_gpt_4v_request",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_non_string_api_keys(patch_openai_client):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test direct request with non-string OpenAI API key (should fail)."
    patch_openai_client("irrelevant")
    with pytest.raises(AttributeError):
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow123",
            openai_api_key=None,
            prompt=[{"role": "user", "content": "Say hello"}],
            gpt_model_version="gpt-4-vision-preview",
            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_0",
            None,
            "test_non_string_api_keys",
            "execute_gpt_4v_request",
            "2_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_invalid_prompt_type_proxied(patch_requests_post):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test proxied request with prompt not being a list (should raise)."
    patch_requests_post(json_data={"choices": [{"message": {"content": "irrelevant"}}]})
    with pytest.raises(TypeError):
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow123",
            openai_api_key="rf_key:account:openai123",
            prompt="not a list",
            gpt_model_version="gpt-4-vision-preview",
            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_0",
            None,
            "test_invalid_prompt_type_proxied",
            "execute_gpt_4v_request",
            "2_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_invalid_prompt_type_direct(patch_openai_client):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test direct request with prompt not being a list (should raise)."
    patch_openai_client("irrelevant")
    with pytest.raises(Exception):
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow123",
            openai_api_key="sk-abc123",
            prompt="not a list",
            gpt_model_version="gpt-4-vision-preview",
            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_0",
            None,
            "test_invalid_prompt_type_direct",
            "execute_gpt_4v_request",
            "2_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_large_prompt_proxied(patch_requests_post):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test proxied request with a large prompt (hundreds of messages)."
    large_prompt = [{"role": "user", "content": f"Message {i}"} for i in range(500)]
    expected_content = "Large prompt processed"
    patch_requests_post(
        json_data={"choices": [{"message": {"content": expected_content}}]}
    )
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow123",
        openai_api_key="rf_key:account:openai123",
        prompt=large_prompt,
        gpt_model_version="gpt-4-vision-preview",
        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_0",
        None,
        "test_large_prompt_proxied",
        "execute_gpt_4v_request",
        "4",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_large_prompt_direct(patch_openai_client):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test direct request with a large prompt (hundreds of messages)."
    large_prompt = [{"role": "user", "content": f"Message {i}"} for i in range(500)]
    expected_content = "Large prompt processed"
    patch_openai_client(expected_content)
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow123",
        openai_api_key="sk-abc123",
        prompt=large_prompt,
        gpt_model_version="gpt-4-vision-preview",
        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_0",
        None,
        "test_large_prompt_direct",
        "execute_gpt_4v_request",
        "4",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_max_tokens_large_value_proxied(patch_requests_post):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test proxied request with a large max_tokens value."
    expected_content = "Large max_tokens processed"
    patch_requests_post(
        json_data={"choices": [{"message": {"content": expected_content}}]}
    )
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow123",
        openai_api_key="rf_key:account:openai123",
        prompt=[{"role": "user", "content": "Test"}],
        gpt_model_version="gpt-4-vision-preview",
        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_0",
        None,
        "test_max_tokens_large_value_proxied",
        "execute_gpt_4v_request",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_max_tokens_large_value_direct(patch_openai_client):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test direct request with a large max_tokens value."
    expected_content = "Large max_tokens processed"
    patch_openai_client(expected_content)
    _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
        roboflow_api_key="rf_key:roboflow123",
        openai_api_key="sk-abc123",
        prompt=[{"role": "user", "content": "Test"}],
        gpt_model_version="gpt-4-vision-preview",
        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_0",
        None,
        "test_max_tokens_large_value_direct",
        "execute_gpt_4v_request",
        "3",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_many_parallel_requests_proxied(patch_requests_post):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test many proxied requests in quick succession (simulate scalability)."
    patch_requests_post(
        json_data={"choices": [{"message": {"content": "Parallel processed"}}]}
    )
    for i in range(50):
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow123",
            openai_api_key="rf_key:account:openai123",
            prompt=[{"role": "user", "content": f"Test {i}"}],
            gpt_model_version="gpt-4-vision-preview",
            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_0",
            None,
            "test_many_parallel_requests_proxied",
            "execute_gpt_4v_request",
            "2_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
        codeflash_output = codeflash_return_value
        result = codeflash_output


def test_many_parallel_requests_direct(patch_openai_client):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "Test many direct requests in quick succession (simulate scalability)."
    patch_openai_client("Parallel processed")
    for i in range(50):
        _call__bound__arguments = inspect.signature(execute_gpt_4v_request).bind(
            roboflow_api_key="rf_key:roboflow123",
            openai_api_key="sk-abc123",
            prompt=[{"role": "user", "content": f"Test {i}"}],
            gpt_model_version="gpt-4-vision-preview",
            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_0",
            None,
            "test_many_parallel_requests_direct",
            "execute_gpt_4v_request",
            "2_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
        codeflash_output = codeflash_return_value
        result = codeflash_output
