from types import GeneratorType

import pytest
import requests
from caikit_nlp_client import HttpClient
from requests.exceptions import SSLError

from .conftest import ConnectionType


def test_generate_text(
    http_client, model_name, prompt, mocker, accept_self_signed_certs
):
    generated_text = http_client.generate_text(model_name, prompt)
    assert isinstance(generated_text, str)
    assert generated_text


def test_generate_text_with_optional_args(
    http_client,
    model_name,
    generated_text_result,
    prompt,
    mocker,
    accept_self_signed_certs,
):
    mock = mocker.spy(requests, "post")

    generated_text = http_client.generate_text(
        model_name, prompt, max_new_tokens=20, min_new_tokens=4
    )

    assert isinstance(generated_text, str)
    assert generated_text
    assert mock.call_args_list[-1].kwargs["json"]["parameters"]["max_new_tokens"] == 20
    assert mock.call_args_list[-1].kwargs["json"]["parameters"]["min_new_tokens"] == 4


def test_timeout_kwarg(
    http_client, model_name, prompt, mocker, accept_self_signed_certs
):
    mock = mocker.spy(requests, "post")

    http_client.generate_text(model_name, prompt)
    assert mock.call_args_list[-1].kwargs["timeout"] == 60.0

    http_client.generate_text(model_name, prompt, timeout=42.0)
    assert mock.call_args_list[-1].kwargs["timeout"] == 42.0


def test_generate_text_with_no_model_id(http_client):
    with pytest.raises(ValueError, match="request must have a model id"):
        http_client.generate_text("", "dummy")


def test_generate_text_stream(
    pytestconfig,
    http_client,
    model_name,
    generated_text_stream_result,
    prompt,
    accept_self_signed_certs,
):
    if not pytestconfig.option.real_caikit:
        pytest.skip(
            reason="stream mocking is broken, see https://github.com/opendatahub-io/caikit-nlp-client/issues/46"
        )

    response = http_client.generate_text_stream(
        model_name,
        prompt,
    )

    assert isinstance(response, GeneratorType)
    response_list = list(response)
    assert response_list
    assert all(isinstance(text, str) for text in response_list)


def test_generate_text_stream_with_optional_args(
    pytestconfig,
    http_client,
    model_name,
    generated_text_stream_result,
    prompt,
    accept_self_signed_certs,
    mocker,
):
    if not pytestconfig.option.real_caikit:
        pytest.skip(
            reason="stream mocking is broken, see https://github.com/opendatahub-io/caikit-nlp-client/issues/46"
        )

    mock = mocker.spy(requests, "post")
    response = http_client.generate_text_stream(
        model_name,
        prompt,
        preserve_input_text=False,
        max_new_tokens=20,
        min_new_tokens=4,
    )

    assert isinstance(response, GeneratorType)
    response_list = list(response)
    assert response_list
    assert all(isinstance(text, str) for text in response_list)
    assert mock.call_args_list[-1].kwargs["json"]["parameters"]["max_new_tokens"] == 20
    assert mock.call_args_list[-1].kwargs["json"]["parameters"]["min_new_tokens"] == 4


def test_get_text_generation_parameters(
    http_client, monkeysession, accept_self_signed_certs
):
    assert http_client.get_text_generation_parameters() == {
        "max_new_tokens": "integer",
        "min_new_tokens": "integer",
        "truncate_input_tokens": "integer",
        "decoding_method": "string",
        "top_k": "integer",
        "top_p": "number",
        "typical_p": "number",
        "temperature": "number",
        "repetition_penalty": "number",
        "max_time": "number",
        "exponential_decay_length_penalty": {
            "start_index": "integer",
            "decay_factor": "number",
        },
        "stop_sequences": "array",
        "seed": "integer",
        "preserve_input_text": "boolean",
    }


@pytest.mark.parametrize("connection_type", [ConnectionType.TLS], indirect=True)
def test_tls_enabled(
    model_name,
    http_server,
    monkeysession,
    ca_cert_file,
    connection_type,
):
    assert connection_type is ConnectionType.TLS, "TLS should be enabled for this test"

    http_client = HttpClient("https://{}:{}".format(*http_server))

    with pytest.raises(SSLError, match=".*CERTIFICATE_VERIFY_FAILED.*"):
        assert http_client.generate_text(model_name, "dummy text")

    # a valid certificate autority should validate the response with no extra args
    with monkeysession.context() as monkeypatch:
        monkeypatch.setenv("REQUESTS_CA_BUNDLE", ca_cert_file)

        assert http_client.generate_text(model_name, "dummy text")


@pytest.mark.parametrize("connection_type", [ConnectionType.MTLS], indirect=True)
def test_client_instantiation(
    ca_cert_file,
    client_key_file,
    client_cert_file,
    connection_type,
):
    # MTLS tests
    with pytest.raises(
        ValueError,
        match="Must provide both ca_cert_path, client_cert_path, client_key_path for "
        "mTLS",
    ):
        HttpClient(
            "https://localhost:8080",
            ca_cert_path=ca_cert_file,
            client_cert_path=client_cert_file,
        )
        HttpClient(
            "https://localhost:8080",
            ca_cert_path=ca_cert_file,
            client_key_path=client_key_file,
        )
        HttpClient(
            "https://localhost:8080",
            client_cert_path=client_cert_file,
            client_key_path=client_key_file,
        )
