import importlib
import json
from unittest import mock

import numpy as np
import openai
import pandas as pd
import pytest
import yaml
from pyspark.sql import SparkSession

import mlflow
import mlflow.pyfunc.scoring_server as pyfunc_scoring_server
from mlflow.models.signature import ModelSignature
from mlflow.models.utils import load_serving_example
from mlflow.types.schema import ColSpec, ParamSchema, ParamSpec, Schema, TensorSpec

from tests.helper_functions import pyfunc_serve_and_score_model
from tests.openai.conftest import is_v1


@pytest.fixture(scope="module")
def spark():
    with SparkSession.builder.master("local[*]").getOrCreate() as s:
        yield s


def chat_completions():
    return openai.chat.completions if is_v1 else openai.ChatCompletion


def completions():
    return openai.completions if is_v1 else openai.Completion


def embeddings():
    return openai.embeddings if is_v1 else openai.Embedding


@pytest.fixture(autouse=True)
def set_envs(monkeypatch, mock_openai):
    monkeypatch.setenv("MLFLOW_TESTING", "true")
    monkeypatch.setenv("OPENAI_API_KEY", "test")
    monkeypatch.setenv("OPENAI_API_BASE", mock_openai)
    if is_v1:
        openai.base_url = mock_openai
    else:
        importlib.reload(openai)


def test_log_model():
    with mlflow.start_run():
        model_info = mlflow.openai.log_model(
            "gpt-4o-mini",
            "chat.completions",
            name="model",
            temperature=0.9,
            messages=[{"role": "system", "content": "You are an MLflow expert."}],
        )

    loaded_model = mlflow.openai.load_model(model_info.model_uri)
    assert loaded_model["model"] == "gpt-4o-mini"
    assert loaded_model["task"] == "chat.completions"
    assert loaded_model["temperature"] == 0.9
    assert loaded_model["messages"] == [{"role": "system", "content": "You are an MLflow expert."}]


def test_chat_single_variable(tmp_path):
    mlflow.openai.save_model(
        model="gpt-4o-mini",
        task=chat_completions(),
        path=tmp_path,
        messages=[{"role": "user", "content": "{x}"}],
    )

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({
        "x": [
            "a",
            "b",
        ]
    })
    expected_output = [
        [{"content": "a", "role": "user"}],
        [{"content": "b", "role": "user"}],
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output

    data = [
        {"x": "a"},
        {"x": "b"},
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output

    data = [
        "a",
        "b",
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output


def test_completion_single_variable(tmp_path):
    mlflow.openai.save_model(
        model="text-davinci-003",
        task=completions(),
        path=tmp_path,
        prompt="Say {text}",
    )

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({
        "x": [
            "this is a test",
            "this is another test",
        ]
    })
    expected_output = ["Say this is a test", "Say this is another test"]
    assert model.predict(data) == expected_output

    data = [
        {"x": "this is a test"},
        {"x": "this is another test"},
    ]
    assert model.predict(data) == expected_output

    data = [
        "this is a test",
        "this is another test",
    ]
    assert model.predict(data) == expected_output


def test_chat_multiple_variables(tmp_path):
    mlflow.openai.save_model(
        model="gpt-4o-mini",
        task=chat_completions(),
        path=tmp_path,
        messages=[{"role": "user", "content": "{x} {y}"}],
    )
    model = mlflow.models.Model.load(tmp_path)
    assert model.signature.inputs.to_dict() == [
        {"name": "x", "type": "string", "required": True},
        {"name": "y", "type": "string", "required": True},
    ]
    assert model.signature.outputs.to_dict() == [
        {"type": "string", "required": True},
    ]

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({
        "x": [
            "a",
            "b",
        ],
        "y": [
            "c",
            "d",
        ],
    })
    expected_output = [
        [{"content": "a c", "role": "user"}],
        [{"content": "b d", "role": "user"}],
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output

    data = [
        {"x": "a", "y": "c"},
        {"x": "b", "y": "d"},
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output


def test_chat_role_content(tmp_path):
    mlflow.openai.save_model(
        model="gpt-4o-mini",
        task=chat_completions(),
        path=tmp_path,
        messages=[{"role": "{role}", "content": "{content}"}],
    )
    model = mlflow.models.Model.load(tmp_path)
    assert model.signature.inputs.to_dict() == [
        {"name": "content", "type": "string", "required": True},
        {"name": "role", "type": "string", "required": True},
    ]
    assert model.signature.outputs.to_dict() == [
        {"type": "string", "required": True},
    ]

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({
        "role": [
            "system",
            "user",
        ],
        "content": [
            "c",
            "d",
        ],
    })
    expected_output = [
        [{"content": "c", "role": "system"}],
        [{"content": "d", "role": "user"}],
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output


def test_completion_multiple_variables(tmp_path):
    mlflow.openai.save_model(
        model="text-davinci-003",
        task=completions(),
        path=tmp_path,
        prompt="Say {x} and {y}",
    )
    model = mlflow.models.Model.load(tmp_path)
    assert model.signature.inputs.to_dict() == [
        {"name": "x", "type": "string", "required": True},
        {"name": "y", "type": "string", "required": True},
    ]
    assert model.signature.outputs.to_dict() == [
        {"type": "string", "required": True},
    ]

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({
        "x": [
            "a",
            "b",
        ],
        "y": [
            "c",
            "d",
        ],
    })
    expected_output = ["Say a and c", "Say b and d"]
    assert model.predict(data) == expected_output

    data = [
        {"x": "a", "y": "c"},
        {"x": "b", "y": "d"},
    ]
    assert model.predict(data) == expected_output


def test_chat_multiple_messages(tmp_path):
    mlflow.openai.save_model(
        model="gpt-4o-mini",
        task=chat_completions(),
        path=tmp_path,
        messages=[
            {"role": "user", "content": "{x}"},
            {"role": "user", "content": "{y}"},
        ],
    )
    model = mlflow.models.Model.load(tmp_path)
    assert model.signature.inputs.to_dict() == [
        {"name": "x", "type": "string", "required": True},
        {"name": "y", "type": "string", "required": True},
    ]
    assert model.signature.outputs.to_dict() == [
        {"type": "string", "required": True},
    ]

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({
        "x": [
            "a",
            "b",
        ],
        "y": [
            "c",
            "d",
        ],
    })
    expected_output = [
        [{"content": "a", "role": "user"}, {"content": "c", "role": "user"}],
        [{"content": "b", "role": "user"}, {"content": "d", "role": "user"}],
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output

    data = [
        {"x": "a", "y": "c"},
        {"x": "b", "y": "d"},
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output


def test_chat_no_variables(tmp_path):
    mlflow.openai.save_model(
        model="gpt-4o-mini",
        task=chat_completions(),
        path=tmp_path,
        messages=[{"role": "user", "content": "a"}],
    )
    model = mlflow.models.Model.load(tmp_path)
    assert model.signature.inputs.to_dict() == [
        {"type": "string", "required": True},
    ]
    assert model.signature.outputs.to_dict() == [
        {"type": "string", "required": True},
    ]

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({
        "content": ["b", "c"],
    })
    expected_output = [
        [{"content": "a", "role": "user"}, {"content": "b", "role": "user"}],
        [{"content": "a", "role": "user"}, {"content": "c", "role": "user"}],
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output

    data = [
        {"content": "b"},
        {"content": "c"},
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output

    data = [
        "b",
        "c",
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output


def test_completion_no_variable(tmp_path):
    mlflow.openai.save_model(
        model="text-davinci-003",
        task=completions(),
        path=tmp_path,
    )

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({
        "x": [
            "this is a test",
            "this is another test",
        ]
    })
    expected_output = ["this is a test", "this is another test"]
    assert model.predict(data) == expected_output

    data = [
        {"x": "this is a test"},
        {"x": "this is another test"},
    ]
    assert model.predict(data) == expected_output

    data = [
        "this is a test",
        "this is another test",
    ]
    assert model.predict(data) == expected_output


def test_chat_no_messages(tmp_path):
    mlflow.openai.save_model(
        model="gpt-4o-mini",
        task=chat_completions(),
        path=tmp_path,
    )
    model = mlflow.models.Model.load(tmp_path)
    assert model.signature.inputs.to_dict() == [
        {"type": "string", "required": True},
    ]
    assert model.signature.outputs.to_dict() == [
        {"type": "string", "required": True},
    ]

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({
        "content": ["b", "c"],
    })
    expected_output = [
        [{"content": "b", "role": "user"}],
        [{"content": "c", "role": "user"}],
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output

    data = [
        {"content": "b"},
        {"content": "c"},
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output

    data = [
        "b",
        "c",
    ]
    assert list(map(json.loads, model.predict(data))) == expected_output


@pytest.mark.parametrize(
    "messages",
    [
        ["a", "b"],
        [{"k": "v"}],
    ],
)
def test_invalid_messages(tmp_path, messages):
    with pytest.raises(
        mlflow.MlflowException,
        match="it must be a list of dictionaries with keys 'role' and 'content'",
    ):
        mlflow.openai.save_model(
            model="gpt-4o-mini",
            task=chat_completions(),
            path=tmp_path,
            messages=messages,
        )


def test_task_argument_accepts_class(tmp_path):
    mlflow.openai.save_model(model="gpt-4o-mini", task=chat_completions(), path=tmp_path)
    loaded_model = mlflow.openai.load_model(tmp_path)
    assert loaded_model["task"] == "chat.completions"


@pytest.mark.skipif(is_v1, reason="Requires OpenAI SDK v0")
def test_model_argument_accepts_retrieved_model(tmp_path):
    model = openai.Model.retrieve("gpt-4o-mini")
    mlflow.openai.save_model(model=model, task=chat_completions(), path=tmp_path)
    loaded_model = mlflow.openai.load_model(tmp_path)
    assert loaded_model["model"] == "gpt-4o-mini"


def test_save_model_with_secret_scope(tmp_path, monkeypatch):
    scope = "test"
    monkeypatch.setenv("MLFLOW_OPENAI_SECRET_SCOPE", scope)
    with (
        mock.patch("mlflow.openai.model.is_in_databricks_runtime", return_value=True),
        mock.patch("mlflow.openai.model.check_databricks_secret_scope_access"),
    ):
        with pytest.warns(FutureWarning, match="MLFLOW_OPENAI_SECRET_SCOPE.+deprecated"):
            mlflow.openai.save_model(model="gpt-4o-mini", task="chat.completions", path=tmp_path)
    with tmp_path.joinpath("openai.yaml").open() as f:
        creds = yaml.safe_load(f)
        assert creds == {
            "OPENAI_API_TYPE": f"{scope}:openai_api_type",
            "OPENAI_API_KEY": f"{scope}:openai_api_key",
            "OPENAI_API_KEY_PATH": f"{scope}:openai_api_key_path",
            "OPENAI_API_BASE": f"{scope}:openai_api_base",
            "OPENAI_BASE_URL": f"{scope}:openai_base_url",
            "OPENAI_ORGANIZATION": f"{scope}:openai_organization",
            "OPENAI_API_VERSION": f"{scope}:openai_api_version",
            "OPENAI_DEPLOYMENT_NAME": f"{scope}:openai_deployment_name",
            "OPENAI_ENGINE": f"{scope}:openai_engine",
        }


def test_spark_udf_chat(tmp_path, spark):
    mlflow.openai.save_model(
        model="gpt-4o-mini",
        task="chat.completions",
        path=tmp_path,
        messages=[
            {"role": "user", "content": "{x} {y}"},
        ],
    )
    udf = mlflow.pyfunc.spark_udf(spark, tmp_path, result_type="string")
    df = spark.createDataFrame(
        [
            ("a", "b"),
            ("c", "d"),
        ],
        ["x", "y"],
    )
    df = df.withColumn("z", udf())
    pdf = df.toPandas()
    assert list(map(json.loads, pdf["z"])) == [
        [{"content": "a b", "role": "user"}],
        [{"content": "c d", "role": "user"}],
    ]


class ChatCompletionModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, params=None):
        completion = chat_completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "What is MLflow?"}],
        )
        return completion.choices[0].message.content


def test_embeddings(tmp_path):
    mlflow.openai.save_model(
        model="text-embedding-ada-002",
        task=embeddings(),
        path=tmp_path,
    )

    model = mlflow.models.Model.load(tmp_path)
    assert model.signature.inputs.to_dict() == [{"type": "string", "required": True}]
    assert model.signature.outputs.to_dict() == [
        {"type": "tensor", "tensor-spec": {"dtype": "float64", "shape": (-1,)}}
    ]

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({"text": ["a", "b"]})
    preds = model.predict(data)
    assert list(map(len, preds)) == [1536, 1536]

    data = pd.DataFrame({"text": ["a"] * 100})
    preds = model.predict(data)
    assert list(map(len, preds)) == [1536] * 100


def test_embeddings_batch_size_azure(tmp_path, monkeypatch):
    monkeypatch.setenv("OPENAI_API_TYPE", "azure")
    monkeypatch.setenv("OPENAI_ENGINE", "test_engine")
    mlflow.openai.save_model(
        model="text-embedding-ada-002",
        task=embeddings(),
        path=tmp_path,
    )
    model = mlflow.pyfunc.load_model(tmp_path)

    assert model._model_impl.api_config.batch_size == 16


def test_embeddings_pyfunc_server_and_score():
    df = pd.DataFrame({"text": ["a", "b"]})
    with mlflow.start_run():
        model_info = mlflow.openai.log_model(
            "text-embedding-ada-002",
            embeddings(),
            name="model",
            input_example=df,
        )
    inference_payload = load_serving_example(model_info.model_uri)
    resp = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    expected = mlflow.pyfunc.load_model(model_info.model_uri).predict(df)
    actual = pd.DataFrame(data=json.loads(resp.content.decode("utf-8")))
    pd.testing.assert_frame_equal(actual, pd.DataFrame({"predictions": expected}))


def test_spark_udf_embeddings(tmp_path, spark):
    mlflow.openai.save_model(
        model="text-embedding-ada-002",
        task=embeddings(),
        path=tmp_path,
    )
    udf = mlflow.pyfunc.spark_udf(spark, tmp_path, result_type="array<double>")
    df = spark.createDataFrame(
        [
            ("a",),
            ("b",),
        ],
        ["x"],
    )
    df = df.withColumn("z", udf("x")).toPandas()
    assert list(map(len, df["z"])) == [1536, 1536]


def test_inference_params(tmp_path):
    mlflow.openai.save_model(
        model="text-embedding-ada-002",
        task=embeddings(),
        path=tmp_path,
        signature=ModelSignature(
            inputs=Schema([ColSpec(type="string", name=None)]),
            outputs=Schema([TensorSpec(type=np.dtype("float64"), shape=(-1,))]),
            params=ParamSchema([ParamSpec(name="batch_size", dtype="long", default=16)]),
        ),
    )

    model_info = mlflow.models.Model.load(tmp_path)
    assert (
        len([p for p in model_info.signature.params if p.name == "batch_size" and p.default == 16])
        == 1
    )

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({"text": ["a", "b"]})
    preds = model.predict(data, params={"batch_size": 5})
    assert list(map(len, preds)) == [1536, 1536]


def test_inference_params_overlap(tmp_path):
    with pytest.raises(mlflow.MlflowException, match=r"any of \['prefix'\] as parameters"):
        mlflow.openai.save_model(
            model="text-davinci-003",
            task=completions(),
            path=tmp_path,
            prefix="Classify the following text's sentiment:",
            signature=ModelSignature(
                inputs=Schema([ColSpec(type="string", name=None)]),
                outputs=Schema([ColSpec(type="string", name=None)]),
                params=ParamSchema([ParamSpec(name="prefix", default=None, dtype="string")]),
            ),
        )


def test_multimodal_messages(tmp_path):
    # Test multimodal content with variable placeholders
    mlflow.openai.save_model(
        model="gpt-4o-mini",
        task=chat_completions(),
        path=tmp_path,
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "{system_prompt}"},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": "data:image/jpeg;base64,{image_base64}",
                            "detail": "low",
                        },
                    },
                ],
            }
        ],
    )

    model = mlflow.models.Model.load(tmp_path)
    assert model.signature.inputs.to_dict() == [
        {"name": "image_base64", "type": "string", "required": True},
        {"name": "system_prompt", "type": "string", "required": True},
    ]
    assert model.signature.outputs.to_dict() == [
        {"type": "string", "required": True},
    ]

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({
        "system_prompt": ["Analyze this image"],
        "image_base64": [
            "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
        ],
    })

    expected_output = [
        [
            {
                "content": [
                    {"type": "text", "text": "Analyze this image"},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": (
                                "data:image/jpeg;base64,"
                                "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
                            ),
                            "detail": "low",
                        },
                    },
                ],
                "role": "user",
            }
        ]
    ]

    assert list(map(json.loads, model.predict(data))) == expected_output


def test_multimodal_messages_no_variables(tmp_path):
    mlflow.openai.save_model(
        model="gpt-4o-mini",
        task=chat_completions(),
        path=tmp_path,
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "What's in this image?"},
                    {
                        "type": "image_url",
                        "image_url": {"url": "data:image/jpeg;base64,abc123", "detail": "low"},
                    },
                ],
            }
        ],
    )

    model = mlflow.models.Model.load(tmp_path)
    # Should add default content variable since no variables found
    assert model.signature.inputs.to_dict() == [
        {"type": "string", "required": True},
    ]

    model = mlflow.pyfunc.load_model(tmp_path)
    data = pd.DataFrame({"content": ["Additional context"]})

    expected_output = [
        [
            {
                "content": [
                    {"type": "text", "text": "What's in this image?"},
                    {
                        "type": "image_url",
                        "image_url": {"url": "data:image/jpeg;base64,abc123", "detail": "low"},
                    },
                ],
                "role": "user",
            },
            {"content": "Additional context", "role": "user"},
        ]
    ]

    assert list(map(json.loads, model.predict(data))) == expected_output