# pylint: disable=line-too-long,useless-suppression
# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------

"""
DESCRIPTION:
    End-to-end scenario combining `.beta.datasets` data generation with an
    evaluation run. The sample:

      1. Creates a `DataGenerationJob` (scenario=EVALUATION, type=simple_qna) that
         synthesizes question/answer pairs from an inline prompt and writes them
         to a new versioned Dataset. Uses `begin_create_generation_job` which returns
         `LROPoller[DataGenerationJobResult]`; `.result()` polls automatically.
      2. Resolves the resulting `DatasetVersion` from the job result.
      3. Creates an OpenAI evaluation (`client.evals.create`) with builtin
         Azure AI evaluators.
      4. Runs the evaluation against the generated dataset by passing the
         dataset's id as the run's `file_id`.
      5. Cleans up the evaluation and the generated dataset.

USAGE:
    python sample_dataset_generation_job_simpleqna_with_prompt_source.py

    Before running the sample:

    pip install "azure-ai-projects>=2.4.0" azure-identity python-dotenv

    Set these environment variables with your own values:
    1) FOUNDRY_PROJECT_ENDPOINT - Required. The Azure AI Project endpoint, as found
       in the overview page of your Microsoft Foundry project.
    2) FOUNDRY_MODEL_NAME - Required. The name of an Azure OpenAI model
       deployment used both to generate the QnA samples and as the judge model
       for builtin evaluators. For `simple_qna` evaluation jobs the deployed
       model must support the Azure OpenAI Responses API. See the supported-model
       list: https://learn.microsoft.com/azure/foundry/openai/how-to/responses?tabs=python-key#model-support
    3) DATASET_NAME - Optional. Name to assign to the generated output dataset.
       Defaults to `dataset-generation-eval-sample`.
    4) POLL_INTERVAL_SECONDS - Optional. Number of seconds to sleep between status
       polls for both the data generation job and the evaluation run. Defaults to 10.
"""

import os
import time
from typing import Union

from dotenv import load_dotenv
from openai.types.eval_create_params import DataSourceConfigCustom
from openai.types.evals.create_eval_completions_run_data_source_param import (
    CreateEvalCompletionsRunDataSourceParam,
    InputMessagesTemplate,
    InputMessagesTemplateTemplateEvalItem,
    SourceFileID,
)
from openai.types.evals.run_create_response import RunCreateResponse
from openai.types.evals.run_retrieve_response import RunRetrieveResponse
from openai.types.responses.response_input_text_param import ResponseInputTextParam

from azure.identity import DefaultAzureCredential
from azure.ai.projects import AIProjectClient
from azure.ai.projects.models import (
    DataGenerationJob,
    DataGenerationJobInputs,
    DataGenerationJobOutputOptions,
    DataGenerationJobScenario,
    DataGenerationModelOptions,
    DatasetDataGenerationJobOutput,
    DatasetVersion,
    PromptDataGenerationJobSource,
    SimpleQnADataGenerationJobOptions,
    TestingCriterionAzureAIEvaluator,
)

load_dotenv()

endpoint = os.environ["FOUNDRY_PROJECT_ENDPOINT"]
model_name = os.environ["FOUNDRY_MODEL_NAME"]
dataset_name = os.environ.get("DATASET_NAME", "dataset-generation-eval-sample")
poll_interval_seconds = int(os.environ.get("POLL_INTERVAL_SECONDS", "10"))


def main() -> None:
    with (
        DefaultAzureCredential() as credential,
        AIProjectClient(endpoint=endpoint, credential=credential) as project_client,
        project_client.get_openai_client() as openai_client,
    ):

        # ------------------------------------------------------------------
        # 1. Generate a QnA evaluation dataset from an inline prompt.
        # ------------------------------------------------------------------
        job = DataGenerationJob(
            inputs=DataGenerationJobInputs(
                name="qna-from-policy-prompt",
                scenario=DataGenerationJobScenario.EVALUATION,
                sources=[
                    PromptDataGenerationJobSource(
                        description="Contoso refund policy",
                        prompt=(
                            "Contoso offers a full refund within 30 days of purchase for any product "
                            "returned in its original condition. After 30 days, store credit may be "
                            "issued at the discretion of customer support. Digital goods are "
                            "non-refundable once downloaded."
                        ),
                    ),
                ],
                options=SimpleQnADataGenerationJobOptions(
                    # Service requires max_samples to be between 15 and 1000.
                    max_samples=15,
                    model_options=DataGenerationModelOptions(model=model_name),
                ),
                output_options=DataGenerationJobOutputOptions(
                    name=dataset_name,
                    description="QnA pairs generated from the Contoso refund policy prompt.",
                    tags={"sample": "dataset-generation-with-evaluation"},
                ),
            ),
        )
        print("Creating data generation job and waiting for completion (polling is handled by the SDK)...")
        job_result = project_client.beta.datasets.begin_create_generation_job(
            job=job,
            polling_interval=poll_interval_seconds,
        ).result()

        # Locate the Dataset output produced by the job.
        output_name: str = ""
        output_version: str = ""
        for output in job_result.outputs or []:
            if isinstance(output, DatasetDataGenerationJobOutput):
                output_name = output.name or ""
                output_version = output.version or ""
                break
        if not output_name or not output_version:
            raise RuntimeError("The data generation job did not produce a dataset output.")

        # Resolve the DatasetVersion so we can use its id as the eval run's file_id.
        dataset: DatasetVersion = project_client.datasets.get(name=output_name, version=output_version)
        print(f"Generated dataset: name=`{dataset.name}` version=`{dataset.version}` id=`{dataset.id}`")

        # ------------------------------------------------------------------
        # 2. Create an evaluation that scores the model's answers to the
        #    generated questions.
        # ------------------------------------------------------------------
        # The `simple_qna` generator produces records with `query` and `ground_truth`
        # fields. The evaluation run below sends each `query` through a model target
        # to produce a response, and the evaluators score that response.
        data_source_config = DataSourceConfigCustom(
            type="custom",
            item_schema={
                "type": "object",
                "properties": {
                    "query": {"type": "string"},
                    "ground_truth": {"type": "string"},
                },
                "required": ["query"],
            },
            include_sample_schema=True,
        )

        testing_criteria = [
            TestingCriterionAzureAIEvaluator(
                type="azure_ai_evaluator",
                name="coherence",
                evaluator_name="builtin.coherence",
                initialization_parameters={"deployment_name": model_name},
                data_mapping={"query": "{{item.query}}", "response": "{{sample.output_text}}"},
            ),
            TestingCriterionAzureAIEvaluator(
                type="azure_ai_evaluator",
                name="fluency",
                evaluator_name="builtin.fluency",
                initialization_parameters={"deployment_name": model_name},
                data_mapping={"response": "{{sample.output_text}}"},
            ),
        ]

        print("Create the evaluation.")
        eval_object = openai_client.evals.create(
            name="generated-qna-evaluation",
            data_source_config=data_source_config,
            testing_criteria=testing_criteria,
        )
        print(f"Evaluation created (id: {eval_object.id}).")

        # ------------------------------------------------------------------
        # 3. Run the evaluation against the generated dataset.
        # ------------------------------------------------------------------
        # `completions` data source: for each item in the dataset, render the
        # input_messages template, send to `model`, and feed the response into
        # the testing criteria as `sample.output_text`.
        print(f"Create an evaluation run that consumes dataset `{dataset.id}`.")
        input_message = InputMessagesTemplate(
            type="template",
            template=[
                InputMessagesTemplateTemplateEvalItem(
                    type="message",
                    role="developer",
                    content=ResponseInputTextParam(
                        type="input_text",
                        text=(
                            "You are a Contoso customer-support assistant. Answer the user's "
                            "question about the Contoso refund policy clearly and concisely."
                        ),
                    ),
                ),
                InputMessagesTemplateTemplateEvalItem(
                    type="message",
                    role="user",
                    content=ResponseInputTextParam(
                        type="input_text",
                        text="{{item.query}}",
                    ),
                ),
            ],
        )
        data_source = CreateEvalCompletionsRunDataSourceParam(
            type="completions",
            source=SourceFileID(type="file_id", id=dataset.id or ""),
            input_messages=input_message,
            model=model_name,
        )
        eval_run: Union[RunCreateResponse, RunRetrieveResponse] = openai_client.evals.runs.create(
            eval_id=eval_object.id,
            name="generated-qna-evaluation-run",
            data_source=data_source,
        )
        print(f"Evaluation run created (id: {eval_run.id}).")

        while eval_run.status not in ("completed", "failed"):
            time.sleep(poll_interval_seconds)
            eval_run = openai_client.evals.runs.retrieve(run_id=eval_run.id, eval_id=eval_object.id)
        print(f"Final eval run status: `{eval_run.status}`.")

        if eval_run.status == "completed":
            print(f"Result counts: {eval_run.result_counts}")
            print(f"Eval run report URL: {eval_run.report_url}")
            output_items = list(openai_client.evals.runs.output_items.list(run_id=eval_run.id, eval_id=eval_object.id))
            print(f"Output items (total: {len(output_items)}):")
            # Print a per-item summary (avoid pprint on the full payload to keep the
            # output ASCII-safe on Windows consoles and easy to scan).
            for idx, item in enumerate(output_items, start=1):
                results = getattr(item, "results", None) or []
                scores = ", ".join(
                    f"{r.get('name', '?')}={r.get('score', 'n/a')} ({r.get('passed', 'n/a')})"
                    for r in results
                    if isinstance(r, dict)
                )
                print(f"  item {idx}: status={item.status} | {scores}")
        else:
            print("Evaluation run did not complete successfully.")

        # ------------------------------------------------------------------
        # 4. Clean up.
        # ------------------------------------------------------------------
        print(f"Delete evaluation `{eval_object.id}`.")
        openai_client.evals.delete(eval_id=eval_object.id)

        print(f"Delete the generated dataset `{dataset.name}` v{dataset.version}.")
        project_client.datasets.delete(name=dataset.name or "", version=dataset.version or "")


if __name__ == "__main__":
    main()
