#
# Copyright 2024 DataRobot, Inc. and its affiliates.
#
# All rights reserved.
#
# DataRobot, Inc.
#
# This is proprietary source code of DataRobot, Inc. and its
# affiliates.
#
# Released under the terms of DataRobot Tool and Utility Agreement.
from __future__ import annotations

from typing import Optional, Union

import trafaret as t

from datarobot.models.api_object import APIObject
from datarobot.models.genai.llm_blueprint import LLMSettingsCommonDict, LLMSettingsCustomModelDict

synthetic_dataset_generation_response_trafaret = t.Dict(
    {
        t.Key("dataset_id"): t.String,
        t.Key("prompt_column_name"): t.String,
        t.Key("response_column_name"): t.String,
    }
).ignore_extra("*")


class SyntheticEvaluationDataset(APIObject):
    """A synthetically generated evaluation dataset for LLMs.

    Attributes
    ----------
    dataset_id (str): The ID of the dataset.
    prompt_column_name (str): The name of the prompt column in the dataset.
    response_column_name (str): The name of the response column in the dataset.
    """

    _path = "api/v2/genai/syntheticEvaluationDatasetGenerations/"

    _converter = synthetic_dataset_generation_response_trafaret

    def __init__(self, dataset_id: str, prompt_column_name: str, response_column_name: str) -> None:
        self.dataset_id = dataset_id
        self.prompt_column_name = prompt_column_name
        self.response_column_name = response_column_name

    @classmethod
    def create(
        cls,
        llm_id: str,
        vector_database_id: str,
        llm_settings: Optional[Union[LLMSettingsCommonDict, LLMSettingsCustomModelDict]] = None,
        dataset_name: Optional[str] = None,
        language: Optional[str] = None,
    ) -> SyntheticEvaluationDataset:
        """Create a synthetic evaluation dataset generation job.  This will
        create a synthetic dataset to be used for evaluation of a language model.

        Parameters
        ----------
        llm_id (str): The ID of the language model that will be used to generate the dataset.
        llm_settings (Dict[Optional[str][Union[bool, int, float, str]]]): The settings to use for the language
                model used for dataset generation.
        vector_database_id (str): The ID of the vector database to use for dataset generation.
        dataset_name (str): The name to use for the dataset, otherwise this will be autogenerated.
        language (str): The language to use for dataset generation.  English will be used if this is not specified.

        Returns
        -------
        SyntheticEvaluationDataset: Reference to the synthetic evaluation dataset that was created.

        """
        url = f"{cls._client.domain}/{cls._path}"
        payload = {
            "llm_id": llm_id,
            "llm_settings": llm_settings,
            "vector_database_id": vector_database_id,
            "dataset_name": dataset_name,
            "language": language,
        }
        response_data = cls._client.post(url, payload)
        return cls.from_server_data(response_data.json())

    def __repr__(self) -> str:
        return (
            f"<SyntheticEvaluationDataset {self.dataset_id}: {self.prompt_column_name} "
            "-> {self.response_column_name}>"
        )
